blob: 3809dc3843c1420fd9ac5c5f45bc8702ea4c308f [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"
Victor Stinner384621c2020-06-22 17:27:35 +02004#include "pycore_tuple.h" // _PyTuple_ITEMS()
Victor Stinner4a21e572020-04-15 02:35:41 +02005#include <stddef.h> // offsetof()
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{
Hai Shi47a23fc2020-06-07 20:05:36 +08001002 Py_VISIT(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001003 Py_VISIT(lz->saved);
1004 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001005}
1006
1007static PyObject *
1008cycle_next(cycleobject *lz)
1009{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001010 PyObject *item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001011
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001012 if (lz->it != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001013 item = PyIter_Next(lz->it);
1014 if (item != NULL) {
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001015 if (lz->firstpass)
1016 return item;
1017 if (PyList_Append(lz->saved, item)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001018 Py_DECREF(item);
1019 return NULL;
1020 }
1021 return item;
1022 }
Raymond Hettinger98958fe2015-08-15 15:09:30 -07001023 /* Note: StopIteration is already cleared by PyIter_Next() */
1024 if (PyErr_Occurred())
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001025 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001026 Py_CLEAR(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001027 }
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001028 if (PyList_GET_SIZE(lz->saved) == 0)
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001029 return NULL;
1030 item = PyList_GET_ITEM(lz->saved, lz->index);
1031 lz->index++;
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001032 if (lz->index >= PyList_GET_SIZE(lz->saved))
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001033 lz->index = 0;
1034 Py_INCREF(item);
1035 return item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001036}
1037
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001038static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301039cycle_reduce(cycleobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001040{
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001041 /* Create a new cycle with the iterator tuple, then set the saved state */
1042 if (lz->it == NULL) {
1043 PyObject *it = PyObject_GetIter(lz->saved);
1044 if (it == NULL)
1045 return NULL;
1046 if (lz->index != 0) {
1047 _Py_IDENTIFIER(__setstate__);
1048 PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
1049 "n", lz->index);
1050 if (res == NULL) {
1051 Py_DECREF(it);
1052 return NULL;
1053 }
1054 Py_DECREF(res);
1055 }
Serhiy Storchaka1f21eaa2019-09-01 12:16:51 +03001056 return Py_BuildValue("O(N)(OO)", Py_TYPE(lz), it, lz->saved, Py_True);
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001057 }
Serhiy Storchaka1f21eaa2019-09-01 12:16:51 +03001058 return Py_BuildValue("O(O)(OO)", Py_TYPE(lz), lz->it, lz->saved,
1059 lz->firstpass ? Py_True : Py_False);
Raymond Hettingerc39786d2015-08-15 15:16:12 -07001060}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001061
1062static PyObject *
1063cycle_setstate(cycleobject *lz, PyObject *state)
1064{
1065 PyObject *saved=NULL;
1066 int firstpass;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001067 if (!PyTuple_Check(state)) {
1068 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001069 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001070 }
1071 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
1072 return NULL;
1073 }
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001074 Py_INCREF(saved);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001075 Py_XSETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001076 lz->firstpass = firstpass != 0;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001077 lz->index = 0;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001078 Py_RETURN_NONE;
1079}
1080
1081static PyMethodDef cycle_methods[] = {
1082 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
1083 reduce_doc},
1084 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
1085 setstate_doc},
1086 {NULL, NULL} /* sentinel */
1087};
1088
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001089static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001090 PyVarObject_HEAD_INIT(NULL, 0)
1091 "itertools.cycle", /* tp_name */
1092 sizeof(cycleobject), /* tp_basicsize */
1093 0, /* tp_itemsize */
1094 /* methods */
1095 (destructor)cycle_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001096 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001097 0, /* tp_getattr */
1098 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001099 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001100 0, /* tp_repr */
1101 0, /* tp_as_number */
1102 0, /* tp_as_sequence */
1103 0, /* tp_as_mapping */
1104 0, /* tp_hash */
1105 0, /* tp_call */
1106 0, /* tp_str */
1107 PyObject_GenericGetAttr, /* tp_getattro */
1108 0, /* tp_setattro */
1109 0, /* tp_as_buffer */
1110 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1111 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001112 itertools_cycle__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001113 (traverseproc)cycle_traverse, /* tp_traverse */
1114 0, /* tp_clear */
1115 0, /* tp_richcompare */
1116 0, /* tp_weaklistoffset */
1117 PyObject_SelfIter, /* tp_iter */
1118 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001119 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001120 0, /* tp_members */
1121 0, /* tp_getset */
1122 0, /* tp_base */
1123 0, /* tp_dict */
1124 0, /* tp_descr_get */
1125 0, /* tp_descr_set */
1126 0, /* tp_dictoffset */
1127 0, /* tp_init */
1128 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001129 itertools_cycle, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001130 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001131};
1132
1133
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001134/* dropwhile object **********************************************************/
1135
1136typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001137 PyObject_HEAD
1138 PyObject *func;
1139 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001140 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001141} dropwhileobject;
1142
Tal Einatc4bccd32018-09-12 00:49:13 +03001143/*[clinic input]
1144@classmethod
1145itertools.dropwhile.__new__
1146 predicate as func: object
1147 iterable as seq: object
1148 /
1149Drop items from the iterable while predicate(item) is true.
1150
1151Afterwards, return every element until the iterable is exhausted.
1152[clinic start generated code]*/
1153
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001154static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001155itertools_dropwhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1156/*[clinic end generated code: output=92f9d0d89af149e4 input=d39737147c9f0a26]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001157{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001158 PyObject *it;
1159 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001160
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001161 /* Get iterator. */
1162 it = PyObject_GetIter(seq);
1163 if (it == NULL)
1164 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001165
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001166 /* create dropwhileobject structure */
1167 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1168 if (lz == NULL) {
1169 Py_DECREF(it);
1170 return NULL;
1171 }
1172 Py_INCREF(func);
1173 lz->func = func;
1174 lz->it = it;
1175 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001176
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001177 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001178}
1179
1180static void
1181dropwhile_dealloc(dropwhileobject *lz)
1182{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001183 PyObject_GC_UnTrack(lz);
1184 Py_XDECREF(lz->func);
1185 Py_XDECREF(lz->it);
1186 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001187}
1188
1189static int
1190dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1191{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001192 Py_VISIT(lz->it);
1193 Py_VISIT(lz->func);
1194 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001195}
1196
1197static PyObject *
1198dropwhile_next(dropwhileobject *lz)
1199{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001200 PyObject *item, *good;
1201 PyObject *it = lz->it;
1202 long ok;
1203 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001205 iternext = *Py_TYPE(it)->tp_iternext;
1206 for (;;) {
1207 item = iternext(it);
1208 if (item == NULL)
1209 return NULL;
1210 if (lz->start == 1)
1211 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001212
Petr Viktorinffd97532020-02-11 17:46:57 +01001213 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001214 if (good == NULL) {
1215 Py_DECREF(item);
1216 return NULL;
1217 }
1218 ok = PyObject_IsTrue(good);
1219 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001220 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001221 lz->start = 1;
1222 return item;
1223 }
1224 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001225 if (ok < 0)
1226 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001227 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001228}
1229
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001230static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301231dropwhile_reduce(dropwhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001232{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001233 return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->start);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001234}
1235
1236static PyObject *
1237dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1238{
1239 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001240 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001241 return NULL;
1242 lz->start = start;
1243 Py_RETURN_NONE;
1244}
1245
1246static PyMethodDef dropwhile_methods[] = {
1247 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1248 reduce_doc},
1249 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1250 setstate_doc},
1251 {NULL, NULL} /* sentinel */
1252};
1253
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001254static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001255 PyVarObject_HEAD_INIT(NULL, 0)
1256 "itertools.dropwhile", /* tp_name */
1257 sizeof(dropwhileobject), /* tp_basicsize */
1258 0, /* tp_itemsize */
1259 /* methods */
1260 (destructor)dropwhile_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001261 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001262 0, /* tp_getattr */
1263 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001264 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001265 0, /* tp_repr */
1266 0, /* tp_as_number */
1267 0, /* tp_as_sequence */
1268 0, /* tp_as_mapping */
1269 0, /* tp_hash */
1270 0, /* tp_call */
1271 0, /* tp_str */
1272 PyObject_GenericGetAttr, /* tp_getattro */
1273 0, /* tp_setattro */
1274 0, /* tp_as_buffer */
1275 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1276 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001277 itertools_dropwhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001278 (traverseproc)dropwhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001279 0, /* tp_clear */
1280 0, /* tp_richcompare */
1281 0, /* tp_weaklistoffset */
1282 PyObject_SelfIter, /* tp_iter */
1283 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001284 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001285 0, /* tp_members */
1286 0, /* tp_getset */
1287 0, /* tp_base */
1288 0, /* tp_dict */
1289 0, /* tp_descr_get */
1290 0, /* tp_descr_set */
1291 0, /* tp_dictoffset */
1292 0, /* tp_init */
1293 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001294 itertools_dropwhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001295 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001296};
1297
1298
1299/* takewhile object **********************************************************/
1300
1301typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001302 PyObject_HEAD
1303 PyObject *func;
1304 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001305 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001306} takewhileobject;
1307
Tal Einatc4bccd32018-09-12 00:49:13 +03001308/*[clinic input]
1309@classmethod
1310itertools.takewhile.__new__
1311 predicate as func: object
1312 iterable as seq: object
1313 /
1314Return successive entries from an iterable as long as the predicate evaluates to true for each entry.
1315[clinic start generated code]*/
1316
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001317static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001318itertools_takewhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1319/*[clinic end generated code: output=bb179ea7864e2ef6 input=ba5255f7519aa119]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001320{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001321 PyObject *it;
1322 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001323
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001324 /* Get iterator. */
1325 it = PyObject_GetIter(seq);
1326 if (it == NULL)
1327 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001328
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001329 /* create takewhileobject structure */
1330 lz = (takewhileobject *)type->tp_alloc(type, 0);
1331 if (lz == NULL) {
1332 Py_DECREF(it);
1333 return NULL;
1334 }
1335 Py_INCREF(func);
1336 lz->func = func;
1337 lz->it = it;
1338 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001339
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001340 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001341}
1342
1343static void
1344takewhile_dealloc(takewhileobject *lz)
1345{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001346 PyObject_GC_UnTrack(lz);
1347 Py_XDECREF(lz->func);
1348 Py_XDECREF(lz->it);
1349 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001350}
1351
1352static int
1353takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1354{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001355 Py_VISIT(lz->it);
1356 Py_VISIT(lz->func);
1357 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001358}
1359
1360static PyObject *
1361takewhile_next(takewhileobject *lz)
1362{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001363 PyObject *item, *good;
1364 PyObject *it = lz->it;
1365 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001366
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001367 if (lz->stop == 1)
1368 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001369
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001370 item = (*Py_TYPE(it)->tp_iternext)(it);
1371 if (item == NULL)
1372 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001373
Petr Viktorinffd97532020-02-11 17:46:57 +01001374 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001375 if (good == NULL) {
1376 Py_DECREF(item);
1377 return NULL;
1378 }
1379 ok = PyObject_IsTrue(good);
1380 Py_DECREF(good);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001381 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001382 return item;
1383 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001384 if (ok == 0)
1385 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001386 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001387}
1388
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001389static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301390takewhile_reduce(takewhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001391{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001392 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 +00001393}
1394
1395static PyObject *
1396takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1397{
1398 int stop = PyObject_IsTrue(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001399
Antoine Pitrou721738f2012-08-15 23:20:39 +02001400 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001401 return NULL;
1402 lz->stop = stop;
1403 Py_RETURN_NONE;
1404}
1405
1406static PyMethodDef takewhile_reduce_methods[] = {
1407 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1408 reduce_doc},
1409 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1410 setstate_doc},
1411 {NULL, NULL} /* sentinel */
1412};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001413
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001414static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001415 PyVarObject_HEAD_INIT(NULL, 0)
1416 "itertools.takewhile", /* tp_name */
1417 sizeof(takewhileobject), /* tp_basicsize */
1418 0, /* tp_itemsize */
1419 /* methods */
1420 (destructor)takewhile_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001421 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001422 0, /* tp_getattr */
1423 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001424 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001425 0, /* tp_repr */
1426 0, /* tp_as_number */
1427 0, /* tp_as_sequence */
1428 0, /* tp_as_mapping */
1429 0, /* tp_hash */
1430 0, /* tp_call */
1431 0, /* tp_str */
1432 PyObject_GenericGetAttr, /* tp_getattro */
1433 0, /* tp_setattro */
1434 0, /* tp_as_buffer */
1435 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1436 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001437 itertools_takewhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001438 (traverseproc)takewhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001439 0, /* tp_clear */
1440 0, /* tp_richcompare */
1441 0, /* tp_weaklistoffset */
1442 PyObject_SelfIter, /* tp_iter */
1443 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001444 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001445 0, /* tp_members */
1446 0, /* tp_getset */
1447 0, /* tp_base */
1448 0, /* tp_dict */
1449 0, /* tp_descr_get */
1450 0, /* tp_descr_set */
1451 0, /* tp_dictoffset */
1452 0, /* tp_init */
1453 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001454 itertools_takewhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001455 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001456};
1457
1458
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001459/* islice object *************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001460
1461typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001462 PyObject_HEAD
1463 PyObject *it;
1464 Py_ssize_t next;
1465 Py_ssize_t stop;
1466 Py_ssize_t step;
1467 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001468} isliceobject;
1469
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001470static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001471
1472static PyObject *
1473islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1474{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001475 PyObject *seq;
1476 Py_ssize_t start=0, stop=-1, step=1;
1477 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1478 Py_ssize_t numargs;
1479 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001480
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001481 if (type == &islice_type && !_PyArg_NoKeywords("islice", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001482 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001483
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001484 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1485 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001486
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001487 numargs = PyTuple_Size(args);
1488 if (numargs == 2) {
1489 if (a1 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001490 stop = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001491 if (stop == -1) {
1492 if (PyErr_Occurred())
1493 PyErr_Clear();
1494 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001495 "Stop argument for islice() must be None or "
1496 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001497 return NULL;
1498 }
1499 }
1500 } else {
1501 if (a1 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001502 start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001503 if (start == -1 && PyErr_Occurred())
1504 PyErr_Clear();
1505 if (a2 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001506 stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001507 if (stop == -1) {
1508 if (PyErr_Occurred())
1509 PyErr_Clear();
1510 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001511 "Stop argument for islice() must be None or "
1512 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001513 return NULL;
1514 }
1515 }
1516 }
1517 if (start<0 || stop<-1) {
1518 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001519 "Indices for islice() must be None or "
1520 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001521 return NULL;
1522 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001523
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001524 if (a3 != NULL) {
1525 if (a3 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001526 step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001527 if (step == -1 && PyErr_Occurred())
1528 PyErr_Clear();
1529 }
1530 if (step<1) {
1531 PyErr_SetString(PyExc_ValueError,
1532 "Step for islice() must be a positive integer or None.");
1533 return NULL;
1534 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001535
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001536 /* Get iterator. */
1537 it = PyObject_GetIter(seq);
1538 if (it == NULL)
1539 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001540
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001541 /* create isliceobject structure */
1542 lz = (isliceobject *)type->tp_alloc(type, 0);
1543 if (lz == NULL) {
1544 Py_DECREF(it);
1545 return NULL;
1546 }
1547 lz->it = it;
1548 lz->next = start;
1549 lz->stop = stop;
1550 lz->step = step;
1551 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001552
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001553 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001554}
1555
1556static void
1557islice_dealloc(isliceobject *lz)
1558{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001559 PyObject_GC_UnTrack(lz);
1560 Py_XDECREF(lz->it);
1561 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001562}
1563
1564static int
1565islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1566{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001567 Py_VISIT(lz->it);
1568 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001569}
1570
1571static PyObject *
1572islice_next(isliceobject *lz)
1573{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001574 PyObject *item;
1575 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001576 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001577 Py_ssize_t oldnext;
1578 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001579
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001580 if (it == NULL)
1581 return NULL;
1582
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001583 iternext = *Py_TYPE(it)->tp_iternext;
1584 while (lz->cnt < lz->next) {
1585 item = iternext(it);
1586 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001587 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001588 Py_DECREF(item);
1589 lz->cnt++;
1590 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001591 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001592 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001593 item = iternext(it);
1594 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001595 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001596 lz->cnt++;
1597 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001598 /* The (size_t) cast below avoids the danger of undefined
1599 behaviour from signed integer overflow. */
1600 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001601 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1602 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001603 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001604
1605empty:
1606 Py_CLEAR(lz->it);
1607 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001608}
1609
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001610static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301611islice_reduce(isliceobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001612{
1613 /* When unpickled, generate a new object with the same bounds,
1614 * then 'setstate' with the next and count
1615 */
1616 PyObject *stop;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001617
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001618 if (lz->it == NULL) {
1619 PyObject *empty_list;
1620 PyObject *empty_it;
1621 empty_list = PyList_New(0);
1622 if (empty_list == NULL)
1623 return NULL;
1624 empty_it = PyObject_GetIter(empty_list);
1625 Py_DECREF(empty_list);
1626 if (empty_it == NULL)
1627 return NULL;
1628 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1629 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001630 if (lz->stop == -1) {
1631 stop = Py_None;
1632 Py_INCREF(stop);
1633 } else {
1634 stop = PyLong_FromSsize_t(lz->stop);
1635 if (stop == NULL)
1636 return NULL;
1637 }
1638 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1639 lz->it, lz->next, stop, lz->step,
1640 lz->cnt);
1641}
1642
1643static PyObject *
1644islice_setstate(isliceobject *lz, PyObject *state)
1645{
1646 Py_ssize_t cnt = PyLong_AsSsize_t(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001647
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001648 if (cnt == -1 && PyErr_Occurred())
1649 return NULL;
1650 lz->cnt = cnt;
1651 Py_RETURN_NONE;
1652}
1653
1654static PyMethodDef islice_methods[] = {
1655 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1656 reduce_doc},
1657 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1658 setstate_doc},
1659 {NULL, NULL} /* sentinel */
1660};
1661
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001662PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001663"islice(iterable, stop) --> islice object\n\
1664islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001665\n\
1666Return an iterator whose next() method returns selected values from an\n\
1667iterable. If start is specified, will skip all preceding elements;\n\
1668otherwise, start defaults to zero. Step defaults to one. If\n\
oldkaa0735f2018-02-02 16:52:55 +08001669specified as another value, step determines how many values are\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001670skipped between successive calls. Works like a slice() on a list\n\
1671but returns an iterator.");
1672
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001673static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001674 PyVarObject_HEAD_INIT(NULL, 0)
1675 "itertools.islice", /* tp_name */
1676 sizeof(isliceobject), /* tp_basicsize */
1677 0, /* tp_itemsize */
1678 /* methods */
1679 (destructor)islice_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001680 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001681 0, /* tp_getattr */
1682 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001683 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001684 0, /* tp_repr */
1685 0, /* tp_as_number */
1686 0, /* tp_as_sequence */
1687 0, /* tp_as_mapping */
1688 0, /* tp_hash */
1689 0, /* tp_call */
1690 0, /* tp_str */
1691 PyObject_GenericGetAttr, /* tp_getattro */
1692 0, /* tp_setattro */
1693 0, /* tp_as_buffer */
1694 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1695 Py_TPFLAGS_BASETYPE, /* tp_flags */
1696 islice_doc, /* tp_doc */
1697 (traverseproc)islice_traverse, /* tp_traverse */
1698 0, /* tp_clear */
1699 0, /* tp_richcompare */
1700 0, /* tp_weaklistoffset */
1701 PyObject_SelfIter, /* tp_iter */
1702 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001703 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001704 0, /* tp_members */
1705 0, /* tp_getset */
1706 0, /* tp_base */
1707 0, /* tp_dict */
1708 0, /* tp_descr_get */
1709 0, /* tp_descr_set */
1710 0, /* tp_dictoffset */
1711 0, /* tp_init */
1712 0, /* tp_alloc */
1713 islice_new, /* tp_new */
1714 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001715};
1716
1717
1718/* starmap object ************************************************************/
1719
1720typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001721 PyObject_HEAD
1722 PyObject *func;
1723 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001724} starmapobject;
1725
Tal Einatc4bccd32018-09-12 00:49:13 +03001726/*[clinic input]
1727@classmethod
1728itertools.starmap.__new__
1729 function as func: object
1730 iterable as seq: object
1731 /
1732Return an iterator whose values are returned from the function evaluated with an argument tuple taken from the given sequence.
1733[clinic start generated code]*/
1734
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001735static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001736itertools_starmap_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1737/*[clinic end generated code: output=79eeb81d452c6e8d input=844766df6a0d4dad]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001738{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001739 PyObject *it;
1740 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001741
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001742 /* Get iterator. */
1743 it = PyObject_GetIter(seq);
1744 if (it == NULL)
1745 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001746
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001747 /* create starmapobject structure */
1748 lz = (starmapobject *)type->tp_alloc(type, 0);
1749 if (lz == NULL) {
1750 Py_DECREF(it);
1751 return NULL;
1752 }
1753 Py_INCREF(func);
1754 lz->func = func;
1755 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001756
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001757 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001758}
1759
1760static void
1761starmap_dealloc(starmapobject *lz)
1762{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001763 PyObject_GC_UnTrack(lz);
1764 Py_XDECREF(lz->func);
1765 Py_XDECREF(lz->it);
1766 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001767}
1768
1769static int
1770starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1771{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001772 Py_VISIT(lz->it);
1773 Py_VISIT(lz->func);
1774 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001775}
1776
1777static PyObject *
1778starmap_next(starmapobject *lz)
1779{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001780 PyObject *args;
1781 PyObject *result;
1782 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001783
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001784 args = (*Py_TYPE(it)->tp_iternext)(it);
1785 if (args == NULL)
1786 return NULL;
1787 if (!PyTuple_CheckExact(args)) {
1788 PyObject *newargs = PySequence_Tuple(args);
1789 Py_DECREF(args);
1790 if (newargs == NULL)
1791 return NULL;
1792 args = newargs;
1793 }
1794 result = PyObject_Call(lz->func, args, NULL);
1795 Py_DECREF(args);
1796 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001797}
1798
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001799static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301800starmap_reduce(starmapobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001801{
1802 /* Just pickle the iterator */
1803 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1804}
1805
1806static PyMethodDef starmap_methods[] = {
1807 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1808 reduce_doc},
1809 {NULL, NULL} /* sentinel */
1810};
1811
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001812static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001813 PyVarObject_HEAD_INIT(NULL, 0)
1814 "itertools.starmap", /* tp_name */
1815 sizeof(starmapobject), /* tp_basicsize */
1816 0, /* tp_itemsize */
1817 /* methods */
1818 (destructor)starmap_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001819 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001820 0, /* tp_getattr */
1821 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001822 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001823 0, /* tp_repr */
1824 0, /* tp_as_number */
1825 0, /* tp_as_sequence */
1826 0, /* tp_as_mapping */
1827 0, /* tp_hash */
1828 0, /* tp_call */
1829 0, /* tp_str */
1830 PyObject_GenericGetAttr, /* tp_getattro */
1831 0, /* tp_setattro */
1832 0, /* tp_as_buffer */
1833 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1834 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001835 itertools_starmap__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001836 (traverseproc)starmap_traverse, /* tp_traverse */
1837 0, /* tp_clear */
1838 0, /* tp_richcompare */
1839 0, /* tp_weaklistoffset */
1840 PyObject_SelfIter, /* tp_iter */
1841 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001842 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001843 0, /* tp_members */
1844 0, /* tp_getset */
1845 0, /* tp_base */
1846 0, /* tp_dict */
1847 0, /* tp_descr_get */
1848 0, /* tp_descr_set */
1849 0, /* tp_dictoffset */
1850 0, /* tp_init */
1851 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001852 itertools_starmap, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001853 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001854};
1855
1856
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001857/* chain object **************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001858
1859typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001860 PyObject_HEAD
1861 PyObject *source; /* Iterator over input iterables */
1862 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001863} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001864
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001865static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001866
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001867static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001868chain_new_internal(PyTypeObject *type, PyObject *source)
1869{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001870 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001871
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001872 lz = (chainobject *)type->tp_alloc(type, 0);
1873 if (lz == NULL) {
1874 Py_DECREF(source);
1875 return NULL;
1876 }
1877
1878 lz->source = source;
1879 lz->active = NULL;
1880 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001881}
1882
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001883static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001884chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001885{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001886 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001887
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001888 if (type == &chain_type && !_PyArg_NoKeywords("chain", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001889 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001890
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001891 source = PyObject_GetIter(args);
1892 if (source == NULL)
1893 return NULL;
1894
1895 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001896}
1897
Tal Einatc4bccd32018-09-12 00:49:13 +03001898/*[clinic input]
1899@classmethod
1900itertools.chain.from_iterable
1901 iterable as arg: object
1902 /
1903Alternative chain() constructor taking a single iterable argument that evaluates lazily.
1904[clinic start generated code]*/
1905
Christian Heimesf16baeb2008-02-29 14:57:44 +00001906static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001907itertools_chain_from_iterable(PyTypeObject *type, PyObject *arg)
1908/*[clinic end generated code: output=667ae7a7f7b68654 input=72c39e3a2ca3be85]*/
Christian Heimesf16baeb2008-02-29 14:57:44 +00001909{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001910 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001911
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001912 source = PyObject_GetIter(arg);
1913 if (source == NULL)
1914 return NULL;
1915
1916 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001917}
1918
1919static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001920chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001921{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001922 PyObject_GC_UnTrack(lz);
1923 Py_XDECREF(lz->active);
1924 Py_XDECREF(lz->source);
1925 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001926}
1927
Raymond Hettinger2012f172003-02-07 05:32:58 +00001928static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001929chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001930{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001931 Py_VISIT(lz->source);
1932 Py_VISIT(lz->active);
1933 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001934}
1935
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001936static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001937chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001938{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001939 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001940
T. Wouters5466d4a2017-03-30 09:58:35 -07001941 /* lz->source is the iterator of iterables. If it's NULL, we've already
1942 * consumed them all. lz->active is the current iterator. If it's NULL,
1943 * we should grab a new one from lz->source. */
1944 while (lz->source != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001945 if (lz->active == NULL) {
T. Wouters5466d4a2017-03-30 09:58:35 -07001946 PyObject *iterable = PyIter_Next(lz->source);
1947 if (iterable == NULL) {
1948 Py_CLEAR(lz->source);
1949 return NULL; /* no more input sources */
1950 }
1951 lz->active = PyObject_GetIter(iterable);
1952 Py_DECREF(iterable);
1953 if (lz->active == NULL) {
1954 Py_CLEAR(lz->source);
1955 return NULL; /* input not iterable */
1956 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001957 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001958 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
1959 if (item != NULL)
1960 return item;
1961 if (PyErr_Occurred()) {
1962 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1963 PyErr_Clear();
1964 else
1965 return NULL; /* input raised an exception */
1966 }
1967 /* lz->active is consumed, try with the next iterable. */
1968 Py_CLEAR(lz->active);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001969 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001970 /* Everything had been consumed already. */
1971 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001972}
1973
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001974static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301975chain_reduce(chainobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001976{
1977 if (lz->source) {
1978 /* we can't pickle function objects (itertools.from_iterable) so
1979 * we must use setstate to replace the iterable. One day we
1980 * will fix pickling of functions
1981 */
1982 if (lz->active) {
1983 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1984 } else {
1985 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1986 }
1987 } else {
1988 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1989 }
1990 return NULL;
1991}
1992
1993static PyObject *
1994chain_setstate(chainobject *lz, PyObject *state)
1995{
1996 PyObject *source, *active=NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001997
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001998 if (!PyTuple_Check(state)) {
1999 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002000 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002001 }
2002 if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
2003 return NULL;
2004 }
2005 if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
2006 PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
2007 return NULL;
2008 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002009
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002010 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002011 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002012 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002013 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002014 Py_RETURN_NONE;
2015}
2016
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002017PyDoc_STRVAR(chain_doc,
2018"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002019\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002020Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002021first iterable until it is exhausted, then elements from the next\n\
2022iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002023
Christian Heimesf16baeb2008-02-29 14:57:44 +00002024static PyMethodDef chain_methods[] = {
Tal Einatc4bccd32018-09-12 00:49:13 +03002025 ITERTOOLS_CHAIN_FROM_ITERABLE_METHODDEF
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002026 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
2027 reduce_doc},
2028 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
2029 setstate_doc},
Ethan Smitha8403d02020-04-09 20:28:08 -07002030 {"__class_getitem__", (PyCFunction)Py_GenericAlias,
2031 METH_O|METH_CLASS, PyDoc_STR("See PEP 585")},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07002032 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00002033};
2034
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002035static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002036 PyVarObject_HEAD_INIT(NULL, 0)
2037 "itertools.chain", /* tp_name */
2038 sizeof(chainobject), /* tp_basicsize */
2039 0, /* tp_itemsize */
2040 /* methods */
2041 (destructor)chain_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002042 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002043 0, /* tp_getattr */
2044 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002045 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002046 0, /* tp_repr */
2047 0, /* tp_as_number */
2048 0, /* tp_as_sequence */
2049 0, /* tp_as_mapping */
2050 0, /* tp_hash */
2051 0, /* tp_call */
2052 0, /* tp_str */
2053 PyObject_GenericGetAttr, /* tp_getattro */
2054 0, /* tp_setattro */
2055 0, /* tp_as_buffer */
2056 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2057 Py_TPFLAGS_BASETYPE, /* tp_flags */
2058 chain_doc, /* tp_doc */
2059 (traverseproc)chain_traverse, /* tp_traverse */
2060 0, /* tp_clear */
2061 0, /* tp_richcompare */
2062 0, /* tp_weaklistoffset */
2063 PyObject_SelfIter, /* tp_iter */
2064 (iternextfunc)chain_next, /* tp_iternext */
2065 chain_methods, /* tp_methods */
2066 0, /* tp_members */
2067 0, /* tp_getset */
2068 0, /* tp_base */
2069 0, /* tp_dict */
2070 0, /* tp_descr_get */
2071 0, /* tp_descr_set */
2072 0, /* tp_dictoffset */
2073 0, /* tp_init */
2074 0, /* tp_alloc */
2075 chain_new, /* tp_new */
2076 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002077};
2078
2079
Christian Heimesc3f30c42008-02-22 16:37:40 +00002080/* product object ************************************************************/
2081
2082typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002083 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002084 PyObject *pools; /* tuple of pool tuples */
2085 Py_ssize_t *indices; /* one index per pool */
2086 PyObject *result; /* most recently returned result tuple */
2087 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002088} productobject;
2089
2090static PyTypeObject product_type;
2091
2092static PyObject *
2093product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2094{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002095 productobject *lz;
2096 Py_ssize_t nargs, npools, repeat=1;
2097 PyObject *pools = NULL;
2098 Py_ssize_t *indices = NULL;
2099 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002100
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002101 if (kwds != NULL) {
2102 char *kwlist[] = {"repeat", 0};
2103 PyObject *tmpargs = PyTuple_New(0);
2104 if (tmpargs == NULL)
2105 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002106 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2107 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002108 Py_DECREF(tmpargs);
2109 return NULL;
2110 }
2111 Py_DECREF(tmpargs);
2112 if (repeat < 0) {
2113 PyErr_SetString(PyExc_ValueError,
2114 "repeat argument cannot be negative");
2115 return NULL;
2116 }
2117 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002118
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002119 assert(PyTuple_CheckExact(args));
2120 if (repeat == 0) {
2121 nargs = 0;
2122 } else {
2123 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002124 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002125 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2126 return NULL;
2127 }
2128 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002129 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002130
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002131 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002132 if (indices == NULL) {
2133 PyErr_NoMemory();
2134 goto error;
2135 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002136
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002137 pools = PyTuple_New(npools);
2138 if (pools == NULL)
2139 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002140
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002141 for (i=0; i < nargs ; ++i) {
2142 PyObject *item = PyTuple_GET_ITEM(args, i);
2143 PyObject *pool = PySequence_Tuple(item);
2144 if (pool == NULL)
2145 goto error;
2146 PyTuple_SET_ITEM(pools, i, pool);
2147 indices[i] = 0;
2148 }
2149 for ( ; i < npools; ++i) {
2150 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2151 Py_INCREF(pool);
2152 PyTuple_SET_ITEM(pools, i, pool);
2153 indices[i] = 0;
2154 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002155
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002156 /* create productobject structure */
2157 lz = (productobject *)type->tp_alloc(type, 0);
2158 if (lz == NULL)
2159 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002160
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002161 lz->pools = pools;
2162 lz->indices = indices;
2163 lz->result = NULL;
2164 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002165
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002166 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002167
2168error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002169 if (indices != NULL)
2170 PyMem_Free(indices);
2171 Py_XDECREF(pools);
2172 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002173}
2174
2175static void
2176product_dealloc(productobject *lz)
2177{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002178 PyObject_GC_UnTrack(lz);
2179 Py_XDECREF(lz->pools);
2180 Py_XDECREF(lz->result);
2181 if (lz->indices != NULL)
2182 PyMem_Free(lz->indices);
2183 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002184}
2185
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002186static PyObject *
2187product_sizeof(productobject *lz, void *unused)
2188{
2189 Py_ssize_t res;
2190
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002191 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002192 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2193 return PyLong_FromSsize_t(res);
2194}
2195
2196PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2197
Christian Heimesc3f30c42008-02-22 16:37:40 +00002198static int
2199product_traverse(productobject *lz, visitproc visit, void *arg)
2200{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002201 Py_VISIT(lz->pools);
2202 Py_VISIT(lz->result);
2203 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002204}
2205
2206static PyObject *
2207product_next(productobject *lz)
2208{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002209 PyObject *pool;
2210 PyObject *elem;
2211 PyObject *oldelem;
2212 PyObject *pools = lz->pools;
2213 PyObject *result = lz->result;
2214 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2215 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002216
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002217 if (lz->stopped)
2218 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002219
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002220 if (result == NULL) {
2221 /* On the first pass, return an initial tuple filled with the
2222 first element from each pool. */
2223 result = PyTuple_New(npools);
2224 if (result == NULL)
2225 goto empty;
2226 lz->result = result;
2227 for (i=0; i < npools; i++) {
2228 pool = PyTuple_GET_ITEM(pools, i);
2229 if (PyTuple_GET_SIZE(pool) == 0)
2230 goto empty;
2231 elem = PyTuple_GET_ITEM(pool, 0);
2232 Py_INCREF(elem);
2233 PyTuple_SET_ITEM(result, i, elem);
2234 }
2235 } else {
2236 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002237
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002238 /* Copy the previous result tuple or re-use it if available */
2239 if (Py_REFCNT(result) > 1) {
2240 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002241 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002242 if (result == NULL)
2243 goto empty;
2244 lz->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002245 Py_DECREF(old_result);
2246 }
2247 /* Now, we've got the only copy so we can update it in-place */
2248 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002249
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002250 /* Update the pool indices right-to-left. Only advance to the
2251 next pool when the previous one rolls-over */
2252 for (i=npools-1 ; i >= 0 ; i--) {
2253 pool = PyTuple_GET_ITEM(pools, i);
2254 indices[i]++;
2255 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2256 /* Roll-over and advance to next pool */
2257 indices[i] = 0;
2258 elem = PyTuple_GET_ITEM(pool, 0);
2259 Py_INCREF(elem);
2260 oldelem = PyTuple_GET_ITEM(result, i);
2261 PyTuple_SET_ITEM(result, i, elem);
2262 Py_DECREF(oldelem);
2263 } else {
2264 /* No rollover. Just increment and stop here. */
2265 elem = PyTuple_GET_ITEM(pool, indices[i]);
2266 Py_INCREF(elem);
2267 oldelem = PyTuple_GET_ITEM(result, i);
2268 PyTuple_SET_ITEM(result, i, elem);
2269 Py_DECREF(oldelem);
2270 break;
2271 }
2272 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002273
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002274 /* If i is negative, then the indices have all rolled-over
2275 and we're done. */
2276 if (i < 0)
2277 goto empty;
2278 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002279
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002280 Py_INCREF(result);
2281 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002282
2283empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002284 lz->stopped = 1;
2285 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002286}
2287
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002288static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302289product_reduce(productobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002290{
2291 if (lz->stopped) {
2292 return Py_BuildValue("O(())", Py_TYPE(lz));
2293 } else if (lz->result == NULL) {
2294 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2295 } else {
2296 PyObject *indices;
2297 Py_ssize_t n, i;
2298
2299 /* we must pickle the indices use them for setstate, and
2300 * additionally indicate that the iterator has started
2301 */
2302 n = PyTuple_GET_SIZE(lz->pools);
2303 indices = PyTuple_New(n);
2304 if (indices == NULL)
2305 return NULL;
2306 for (i=0; i<n; i++){
2307 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2308 if (!index) {
2309 Py_DECREF(indices);
2310 return NULL;
2311 }
2312 PyTuple_SET_ITEM(indices, i, index);
2313 }
2314 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2315 }
2316}
2317
2318static PyObject *
2319product_setstate(productobject *lz, PyObject *state)
2320{
2321 PyObject *result;
2322 Py_ssize_t n, i;
2323
2324 n = PyTuple_GET_SIZE(lz->pools);
2325 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2326 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2327 return NULL;
2328 }
2329 for (i=0; i<n; i++)
2330 {
2331 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2332 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002333 PyObject* pool;
2334 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002335 if (index < 0 && PyErr_Occurred())
2336 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002337 pool = PyTuple_GET_ITEM(lz->pools, i);
2338 poolsize = PyTuple_GET_SIZE(pool);
2339 if (poolsize == 0) {
2340 lz->stopped = 1;
2341 Py_RETURN_NONE;
2342 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002343 /* clamp the index */
2344 if (index < 0)
2345 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002346 else if (index > poolsize-1)
2347 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002348 lz->indices[i] = index;
2349 }
2350
2351 result = PyTuple_New(n);
2352 if (!result)
2353 return NULL;
2354 for (i=0; i<n; i++) {
2355 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2356 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2357 Py_INCREF(element);
2358 PyTuple_SET_ITEM(result, i, element);
2359 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002360 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002361 Py_RETURN_NONE;
2362}
2363
2364static PyMethodDef product_methods[] = {
2365 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2366 reduce_doc},
2367 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2368 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002369 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2370 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002371 {NULL, NULL} /* sentinel */
2372};
2373
Christian Heimesc3f30c42008-02-22 16:37:40 +00002374PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002375"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002376\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002377Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002378For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2379The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2380cycle in a manner similar to an odometer (with the rightmost element changing\n\
2381on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002382To compute the product of an iterable with itself, specify the number\n\
2383of repetitions with the optional repeat keyword argument. For example,\n\
2384product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002385product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2386product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2387
2388static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002389 PyVarObject_HEAD_INIT(NULL, 0)
2390 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002391 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002392 0, /* tp_itemsize */
2393 /* methods */
2394 (destructor)product_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002395 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002396 0, /* tp_getattr */
2397 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002398 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002399 0, /* tp_repr */
2400 0, /* tp_as_number */
2401 0, /* tp_as_sequence */
2402 0, /* tp_as_mapping */
2403 0, /* tp_hash */
2404 0, /* tp_call */
2405 0, /* tp_str */
2406 PyObject_GenericGetAttr, /* tp_getattro */
2407 0, /* tp_setattro */
2408 0, /* tp_as_buffer */
2409 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2410 Py_TPFLAGS_BASETYPE, /* tp_flags */
2411 product_doc, /* tp_doc */
2412 (traverseproc)product_traverse, /* tp_traverse */
2413 0, /* tp_clear */
2414 0, /* tp_richcompare */
2415 0, /* tp_weaklistoffset */
2416 PyObject_SelfIter, /* tp_iter */
2417 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002418 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002419 0, /* tp_members */
2420 0, /* tp_getset */
2421 0, /* tp_base */
2422 0, /* tp_dict */
2423 0, /* tp_descr_get */
2424 0, /* tp_descr_set */
2425 0, /* tp_dictoffset */
2426 0, /* tp_init */
2427 0, /* tp_alloc */
2428 product_new, /* tp_new */
2429 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002430};
2431
2432
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002433/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002434
2435typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002436 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002437 PyObject *pool; /* input converted to a tuple */
2438 Py_ssize_t *indices; /* one index per result element */
2439 PyObject *result; /* most recently returned result tuple */
2440 Py_ssize_t r; /* size of result tuple */
2441 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002442} combinationsobject;
2443
Tal Einatc4bccd32018-09-12 00:49:13 +03002444
2445/*[clinic input]
2446@classmethod
2447itertools.combinations.__new__
2448 iterable: object
2449 r: Py_ssize_t
2450Return successive r-length combinations of elements in the iterable.
2451
2452combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)
2453[clinic start generated code]*/
2454
Christian Heimes380f7f22008-02-28 11:19:05 +00002455static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002456itertools_combinations_impl(PyTypeObject *type, PyObject *iterable,
2457 Py_ssize_t r)
2458/*[clinic end generated code: output=87a689b39c40039c input=06bede09e3da20f8]*/
Christian Heimes380f7f22008-02-28 11:19:05 +00002459{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002460 combinationsobject *co;
2461 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002462 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002463 Py_ssize_t *indices = NULL;
2464 Py_ssize_t i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002465
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002466 pool = PySequence_Tuple(iterable);
2467 if (pool == NULL)
2468 goto error;
2469 n = PyTuple_GET_SIZE(pool);
2470 if (r < 0) {
2471 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2472 goto error;
2473 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002474
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002475 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002476 if (indices == NULL) {
2477 PyErr_NoMemory();
2478 goto error;
2479 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002480
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002481 for (i=0 ; i<r ; i++)
2482 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002483
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002484 /* create combinationsobject structure */
2485 co = (combinationsobject *)type->tp_alloc(type, 0);
2486 if (co == NULL)
2487 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002488
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002489 co->pool = pool;
2490 co->indices = indices;
2491 co->result = NULL;
2492 co->r = r;
2493 co->stopped = r > n ? 1 : 0;
2494
2495 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002496
2497error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002498 if (indices != NULL)
2499 PyMem_Free(indices);
2500 Py_XDECREF(pool);
2501 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002502}
2503
2504static void
2505combinations_dealloc(combinationsobject *co)
2506{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002507 PyObject_GC_UnTrack(co);
2508 Py_XDECREF(co->pool);
2509 Py_XDECREF(co->result);
2510 if (co->indices != NULL)
2511 PyMem_Free(co->indices);
2512 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002513}
2514
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002515static PyObject *
2516combinations_sizeof(combinationsobject *co, void *unused)
2517{
2518 Py_ssize_t res;
2519
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002520 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002521 res += co->r * sizeof(Py_ssize_t);
2522 return PyLong_FromSsize_t(res);
2523}
2524
Christian Heimes380f7f22008-02-28 11:19:05 +00002525static int
2526combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2527{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002528 Py_VISIT(co->pool);
2529 Py_VISIT(co->result);
2530 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002531}
2532
2533static PyObject *
2534combinations_next(combinationsobject *co)
2535{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002536 PyObject *elem;
2537 PyObject *oldelem;
2538 PyObject *pool = co->pool;
2539 Py_ssize_t *indices = co->indices;
2540 PyObject *result = co->result;
2541 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2542 Py_ssize_t r = co->r;
2543 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002544
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002545 if (co->stopped)
2546 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002547
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002548 if (result == NULL) {
2549 /* On the first pass, initialize result tuple using the indices */
2550 result = PyTuple_New(r);
2551 if (result == NULL)
2552 goto empty;
2553 co->result = result;
2554 for (i=0; i<r ; i++) {
2555 index = indices[i];
2556 elem = PyTuple_GET_ITEM(pool, index);
2557 Py_INCREF(elem);
2558 PyTuple_SET_ITEM(result, i, elem);
2559 }
2560 } else {
2561 /* Copy the previous result tuple or re-use it if available */
2562 if (Py_REFCNT(result) > 1) {
2563 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002564 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002565 if (result == NULL)
2566 goto empty;
2567 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002568 Py_DECREF(old_result);
2569 }
2570 /* Now, we've got the only copy so we can update it in-place
2571 * CPython's empty tuple is a singleton and cached in
2572 * PyTuple's freelist.
2573 */
2574 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002575
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002576 /* Scan indices right-to-left until finding one that is not
2577 at its maximum (i + n - r). */
2578 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2579 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002580
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002581 /* If i is negative, then the indices are all at
2582 their maximum value and we're done. */
2583 if (i < 0)
2584 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002585
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002586 /* Increment the current index which we know is not at its
2587 maximum. Then move back to the right setting each index
2588 to its lowest possible value (one higher than the index
2589 to its left -- this maintains the sort order invariant). */
2590 indices[i]++;
2591 for (j=i+1 ; j<r ; j++)
2592 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002593
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002594 /* Update the result tuple for the new indices
2595 starting with i, the leftmost index that changed */
2596 for ( ; i<r ; i++) {
2597 index = indices[i];
2598 elem = PyTuple_GET_ITEM(pool, index);
2599 Py_INCREF(elem);
2600 oldelem = PyTuple_GET_ITEM(result, i);
2601 PyTuple_SET_ITEM(result, i, elem);
2602 Py_DECREF(oldelem);
2603 }
2604 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002605
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002606 Py_INCREF(result);
2607 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002608
2609empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002610 co->stopped = 1;
2611 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002612}
2613
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002614static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302615combinations_reduce(combinationsobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002616{
2617 if (lz->result == NULL) {
2618 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2619 } else if (lz->stopped) {
2620 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2621 } else {
2622 PyObject *indices;
2623 Py_ssize_t i;
2624
2625 /* we must pickle the indices and use them for setstate */
2626 indices = PyTuple_New(lz->r);
2627 if (!indices)
2628 return NULL;
2629 for (i=0; i<lz->r; i++)
2630 {
2631 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2632 if (!index) {
2633 Py_DECREF(indices);
2634 return NULL;
2635 }
2636 PyTuple_SET_ITEM(indices, i, index);
2637 }
2638
2639 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2640 }
2641}
2642
2643static PyObject *
2644combinations_setstate(combinationsobject *lz, PyObject *state)
2645{
2646 PyObject *result;
2647 Py_ssize_t i;
2648 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2649
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002650 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002651 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2652 return NULL;
2653 }
2654
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002655 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002656 Py_ssize_t max;
2657 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2658 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002659
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002660 if (index == -1 && PyErr_Occurred())
2661 return NULL; /* not an integer */
2662 max = i + n - lz->r;
2663 /* clamp the index (beware of negative max) */
2664 if (index > max)
2665 index = max;
2666 if (index < 0)
2667 index = 0;
2668 lz->indices[i] = index;
2669 }
2670
2671 result = PyTuple_New(lz->r);
2672 if (result == NULL)
2673 return NULL;
2674 for (i=0; i<lz->r; i++) {
2675 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2676 Py_INCREF(element);
2677 PyTuple_SET_ITEM(result, i, element);
2678 }
2679
Serhiy Storchaka48842712016-04-06 09:45:48 +03002680 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002681 Py_RETURN_NONE;
2682}
2683
2684static PyMethodDef combinations_methods[] = {
2685 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2686 reduce_doc},
2687 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2688 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002689 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2690 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002691 {NULL, NULL} /* sentinel */
2692};
2693
Christian Heimes380f7f22008-02-28 11:19:05 +00002694static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002695 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002696 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002697 sizeof(combinationsobject), /* tp_basicsize */
2698 0, /* tp_itemsize */
2699 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002700 (destructor)combinations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002701 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002702 0, /* tp_getattr */
2703 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002704 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002705 0, /* tp_repr */
2706 0, /* tp_as_number */
2707 0, /* tp_as_sequence */
2708 0, /* tp_as_mapping */
2709 0, /* tp_hash */
2710 0, /* tp_call */
2711 0, /* tp_str */
2712 PyObject_GenericGetAttr, /* tp_getattro */
2713 0, /* tp_setattro */
2714 0, /* tp_as_buffer */
2715 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2716 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03002717 itertools_combinations__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002718 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002719 0, /* tp_clear */
2720 0, /* tp_richcompare */
2721 0, /* tp_weaklistoffset */
2722 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002723 (iternextfunc)combinations_next, /* tp_iternext */
2724 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002725 0, /* tp_members */
2726 0, /* tp_getset */
2727 0, /* tp_base */
2728 0, /* tp_dict */
2729 0, /* tp_descr_get */
2730 0, /* tp_descr_set */
2731 0, /* tp_dictoffset */
2732 0, /* tp_init */
2733 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03002734 itertools_combinations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002735 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002736};
2737
2738
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002739/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002740
2741/* Equivalent to:
2742
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002743 def combinations_with_replacement(iterable, r):
2744 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2745 # number items returned: (n+r-1)! / r! / (n-1)!
2746 pool = tuple(iterable)
2747 n = len(pool)
2748 indices = [0] * r
2749 yield tuple(pool[i] for i in indices)
2750 while 1:
2751 for i in reversed(range(r)):
2752 if indices[i] != n - 1:
2753 break
2754 else:
2755 return
2756 indices[i:] = [indices[i] + 1] * (r - i)
2757 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002758
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002759 def combinations_with_replacement2(iterable, r):
2760 'Alternate version that filters from product()'
2761 pool = tuple(iterable)
2762 n = len(pool)
2763 for indices in product(range(n), repeat=r):
2764 if sorted(indices) == list(indices):
2765 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002766*/
2767typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002768 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002769 PyObject *pool; /* input converted to a tuple */
2770 Py_ssize_t *indices; /* one index per result element */
2771 PyObject *result; /* most recently returned result tuple */
2772 Py_ssize_t r; /* size of result tuple */
2773 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002774} cwrobject;
2775
Tal Einatc4bccd32018-09-12 00:49:13 +03002776/*[clinic input]
2777@classmethod
2778itertools.combinations_with_replacement.__new__
2779 iterable: object
2780 r: Py_ssize_t
2781Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats.
2782
2783combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2784[clinic start generated code]*/
2785
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002786static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002787itertools_combinations_with_replacement_impl(PyTypeObject *type,
2788 PyObject *iterable,
2789 Py_ssize_t r)
2790/*[clinic end generated code: output=48b26856d4e659ca input=dc2a8c7ba785fad7]*/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002791{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002792 cwrobject *co;
2793 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002794 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002795 Py_ssize_t *indices = NULL;
2796 Py_ssize_t i;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002797
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002798 pool = PySequence_Tuple(iterable);
2799 if (pool == NULL)
2800 goto error;
2801 n = PyTuple_GET_SIZE(pool);
2802 if (r < 0) {
2803 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2804 goto error;
2805 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002806
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002807 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002808 if (indices == NULL) {
2809 PyErr_NoMemory();
2810 goto error;
2811 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002812
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002813 for (i=0 ; i<r ; i++)
2814 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002815
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002816 /* create cwrobject structure */
2817 co = (cwrobject *)type->tp_alloc(type, 0);
2818 if (co == NULL)
2819 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002820
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002821 co->pool = pool;
2822 co->indices = indices;
2823 co->result = NULL;
2824 co->r = r;
2825 co->stopped = !n && r;
2826
2827 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002828
2829error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002830 if (indices != NULL)
2831 PyMem_Free(indices);
2832 Py_XDECREF(pool);
2833 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002834}
2835
2836static void
2837cwr_dealloc(cwrobject *co)
2838{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002839 PyObject_GC_UnTrack(co);
2840 Py_XDECREF(co->pool);
2841 Py_XDECREF(co->result);
2842 if (co->indices != NULL)
2843 PyMem_Free(co->indices);
2844 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002845}
2846
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002847static PyObject *
2848cwr_sizeof(cwrobject *co, void *unused)
2849{
2850 Py_ssize_t res;
2851
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002852 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002853 res += co->r * sizeof(Py_ssize_t);
2854 return PyLong_FromSsize_t(res);
2855}
2856
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002857static int
2858cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2859{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002860 Py_VISIT(co->pool);
2861 Py_VISIT(co->result);
2862 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002863}
2864
2865static PyObject *
2866cwr_next(cwrobject *co)
2867{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002868 PyObject *elem;
2869 PyObject *oldelem;
2870 PyObject *pool = co->pool;
2871 Py_ssize_t *indices = co->indices;
2872 PyObject *result = co->result;
2873 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2874 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002875 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002876
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002877 if (co->stopped)
2878 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002879
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002880 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002881 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002882 result = PyTuple_New(r);
2883 if (result == NULL)
2884 goto empty;
2885 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002886 if (n > 0) {
2887 elem = PyTuple_GET_ITEM(pool, 0);
2888 for (i=0; i<r ; i++) {
2889 assert(indices[i] == 0);
2890 Py_INCREF(elem);
2891 PyTuple_SET_ITEM(result, i, elem);
2892 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002893 }
2894 } else {
2895 /* Copy the previous result tuple or re-use it if available */
2896 if (Py_REFCNT(result) > 1) {
2897 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002898 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002899 if (result == NULL)
2900 goto empty;
2901 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002902 Py_DECREF(old_result);
2903 }
2904 /* Now, we've got the only copy so we can update it in-place CPython's
2905 empty tuple is a singleton and cached in PyTuple's freelist. */
2906 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002907
Tim Peters9edb1682013-09-03 11:49:31 -05002908 /* Scan indices right-to-left until finding one that is not
2909 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002910 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2911 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002912
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002913 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002914 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002915 if (i < 0)
2916 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002917
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002918 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002919 maximum. Then set all to the right to the same value. */
2920 index = indices[i] + 1;
2921 assert(index < n);
2922 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002923 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002924 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002925 Py_INCREF(elem);
2926 oldelem = PyTuple_GET_ITEM(result, i);
2927 PyTuple_SET_ITEM(result, i, elem);
2928 Py_DECREF(oldelem);
2929 }
2930 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002931
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002932 Py_INCREF(result);
2933 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002934
2935empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002936 co->stopped = 1;
2937 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002938}
2939
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002940static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302941cwr_reduce(cwrobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002942{
2943 if (lz->result == NULL) {
2944 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2945 } else if (lz->stopped) {
2946 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2947 } else {
2948 PyObject *indices;
2949 Py_ssize_t i;
2950
2951 /* we must pickle the indices and use them for setstate */
2952 indices = PyTuple_New(lz->r);
2953 if (!indices)
2954 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002955 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002956 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2957 if (!index) {
2958 Py_DECREF(indices);
2959 return NULL;
2960 }
2961 PyTuple_SET_ITEM(indices, i, index);
2962 }
2963
2964 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2965 }
2966}
2967
2968static PyObject *
2969cwr_setstate(cwrobject *lz, PyObject *state)
2970{
2971 PyObject *result;
2972 Py_ssize_t n, i;
2973
2974 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2975 {
2976 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2977 return NULL;
2978 }
2979
2980 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002981 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002982 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2983 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002984
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002985 if (index < 0 && PyErr_Occurred())
2986 return NULL; /* not an integer */
2987 /* clamp the index */
2988 if (index < 0)
2989 index = 0;
2990 else if (index > n-1)
2991 index = n-1;
2992 lz->indices[i] = index;
2993 }
2994 result = PyTuple_New(lz->r);
2995 if (result == NULL)
2996 return NULL;
2997 for (i=0; i<lz->r; i++) {
2998 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2999 Py_INCREF(element);
3000 PyTuple_SET_ITEM(result, i, element);
3001 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003002 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003003 Py_RETURN_NONE;
3004}
3005
3006static PyMethodDef cwr_methods[] = {
3007 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
3008 reduce_doc},
3009 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
3010 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003011 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
3012 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003013 {NULL, NULL} /* sentinel */
3014};
3015
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003016static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003017 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003018 "itertools.combinations_with_replacement", /* tp_name */
3019 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003020 0, /* tp_itemsize */
3021 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003022 (destructor)cwr_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003023 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003024 0, /* tp_getattr */
3025 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003026 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003027 0, /* tp_repr */
3028 0, /* tp_as_number */
3029 0, /* tp_as_sequence */
3030 0, /* tp_as_mapping */
3031 0, /* tp_hash */
3032 0, /* tp_call */
3033 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003034 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003035 0, /* tp_setattro */
3036 0, /* tp_as_buffer */
3037 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003038 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003039 itertools_combinations_with_replacement__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003040 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003041 0, /* tp_clear */
3042 0, /* tp_richcompare */
3043 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003044 PyObject_SelfIter, /* tp_iter */
3045 (iternextfunc)cwr_next, /* tp_iternext */
3046 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003047 0, /* tp_members */
3048 0, /* tp_getset */
3049 0, /* tp_base */
3050 0, /* tp_dict */
3051 0, /* tp_descr_get */
3052 0, /* tp_descr_set */
3053 0, /* tp_dictoffset */
3054 0, /* tp_init */
3055 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003056 itertools_combinations_with_replacement, /* tp_new */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003057 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003058};
3059
3060
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003061/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003062
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003063def permutations(iterable, r=None):
Sergey0078a0c2019-10-29 08:10:24 +03003064 # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
3065 # permutations(range(3)) --> 012 021 102 120 201 210
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003066 pool = tuple(iterable)
3067 n = len(pool)
3068 r = n if r is None else r
Sergey0078a0c2019-10-29 08:10:24 +03003069 if r > n:
3070 return
3071 indices = list(range(n))
3072 cycles = list(range(n, n-r, -1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003073 yield tuple(pool[i] for i in indices[:r])
3074 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003075 for i in reversed(range(r)):
3076 cycles[i] -= 1
3077 if cycles[i] == 0:
3078 indices[i:] = indices[i+1:] + indices[i:i+1]
3079 cycles[i] = n - i
3080 else:
3081 j = cycles[i]
3082 indices[i], indices[-j] = indices[-j], indices[i]
3083 yield tuple(pool[i] for i in indices[:r])
3084 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003085 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003086 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003087*/
3088
3089typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003090 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003091 PyObject *pool; /* input converted to a tuple */
3092 Py_ssize_t *indices; /* one index per element in the pool */
3093 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3094 PyObject *result; /* most recently returned result tuple */
3095 Py_ssize_t r; /* size of result tuple */
3096 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003097} permutationsobject;
3098
Tal Einatc4bccd32018-09-12 00:49:13 +03003099/*[clinic input]
3100@classmethod
3101itertools.permutations.__new__
3102 iterable: object
3103 r as robj: object = None
3104Return successive r-length permutations of elements in the iterable.
3105
3106permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)
3107[clinic start generated code]*/
3108
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003109static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003110itertools_permutations_impl(PyTypeObject *type, PyObject *iterable,
3111 PyObject *robj)
3112/*[clinic end generated code: output=296a72fa76d620ea input=57d0170a4ac0ec7a]*/
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003113{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003114 permutationsobject *po;
3115 Py_ssize_t n;
3116 Py_ssize_t r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003117 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003118 Py_ssize_t *indices = NULL;
3119 Py_ssize_t *cycles = NULL;
3120 Py_ssize_t i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003121
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003122 pool = PySequence_Tuple(iterable);
3123 if (pool == NULL)
3124 goto error;
3125 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003126
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003127 r = n;
3128 if (robj != Py_None) {
3129 if (!PyLong_Check(robj)) {
3130 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3131 goto error;
3132 }
3133 r = PyLong_AsSsize_t(robj);
3134 if (r == -1 && PyErr_Occurred())
3135 goto error;
3136 }
3137 if (r < 0) {
3138 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3139 goto error;
3140 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003141
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003142 indices = PyMem_New(Py_ssize_t, n);
3143 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003144 if (indices == NULL || cycles == NULL) {
3145 PyErr_NoMemory();
3146 goto error;
3147 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003148
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003149 for (i=0 ; i<n ; i++)
3150 indices[i] = i;
3151 for (i=0 ; i<r ; i++)
3152 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003153
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003154 /* create permutationsobject structure */
3155 po = (permutationsobject *)type->tp_alloc(type, 0);
3156 if (po == NULL)
3157 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003158
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003159 po->pool = pool;
3160 po->indices = indices;
3161 po->cycles = cycles;
3162 po->result = NULL;
3163 po->r = r;
3164 po->stopped = r > n ? 1 : 0;
3165
3166 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003167
3168error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003169 if (indices != NULL)
3170 PyMem_Free(indices);
3171 if (cycles != NULL)
3172 PyMem_Free(cycles);
3173 Py_XDECREF(pool);
3174 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003175}
3176
3177static void
3178permutations_dealloc(permutationsobject *po)
3179{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003180 PyObject_GC_UnTrack(po);
3181 Py_XDECREF(po->pool);
3182 Py_XDECREF(po->result);
3183 PyMem_Free(po->indices);
3184 PyMem_Free(po->cycles);
3185 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003186}
3187
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003188static PyObject *
3189permutations_sizeof(permutationsobject *po, void *unused)
3190{
3191 Py_ssize_t res;
3192
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003193 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003194 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3195 res += po->r * sizeof(Py_ssize_t);
3196 return PyLong_FromSsize_t(res);
3197}
3198
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003199static int
3200permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3201{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003202 Py_VISIT(po->pool);
3203 Py_VISIT(po->result);
3204 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003205}
3206
3207static PyObject *
3208permutations_next(permutationsobject *po)
3209{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003210 PyObject *elem;
3211 PyObject *oldelem;
3212 PyObject *pool = po->pool;
3213 Py_ssize_t *indices = po->indices;
3214 Py_ssize_t *cycles = po->cycles;
3215 PyObject *result = po->result;
3216 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3217 Py_ssize_t r = po->r;
3218 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003219
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003220 if (po->stopped)
3221 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003222
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003223 if (result == NULL) {
3224 /* On the first pass, initialize result tuple using the indices */
3225 result = PyTuple_New(r);
3226 if (result == NULL)
3227 goto empty;
3228 po->result = result;
3229 for (i=0; i<r ; i++) {
3230 index = indices[i];
3231 elem = PyTuple_GET_ITEM(pool, index);
3232 Py_INCREF(elem);
3233 PyTuple_SET_ITEM(result, i, elem);
3234 }
3235 } else {
3236 if (n == 0)
3237 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003238
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003239 /* Copy the previous result tuple or re-use it if available */
3240 if (Py_REFCNT(result) > 1) {
3241 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05003242 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003243 if (result == NULL)
3244 goto empty;
3245 po->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003246 Py_DECREF(old_result);
3247 }
3248 /* Now, we've got the only copy so we can update it in-place */
3249 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003250
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003251 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3252 for (i=r-1 ; i>=0 ; i--) {
3253 cycles[i] -= 1;
3254 if (cycles[i] == 0) {
3255 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3256 index = indices[i];
3257 for (j=i ; j<n-1 ; j++)
3258 indices[j] = indices[j+1];
3259 indices[n-1] = index;
3260 cycles[i] = n - i;
3261 } else {
3262 j = cycles[i];
3263 index = indices[i];
3264 indices[i] = indices[n-j];
3265 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003266
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003267 for (k=i; k<r ; k++) {
3268 /* start with i, the leftmost element that changed */
3269 /* yield tuple(pool[k] for k in indices[:r]) */
3270 index = indices[k];
3271 elem = PyTuple_GET_ITEM(pool, index);
3272 Py_INCREF(elem);
3273 oldelem = PyTuple_GET_ITEM(result, k);
3274 PyTuple_SET_ITEM(result, k, elem);
3275 Py_DECREF(oldelem);
3276 }
3277 break;
3278 }
3279 }
3280 /* If i is negative, then the cycles have all
3281 rolled-over and we're done. */
3282 if (i < 0)
3283 goto empty;
3284 }
3285 Py_INCREF(result);
3286 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003287
3288empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003289 po->stopped = 1;
3290 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003291}
3292
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003293static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303294permutations_reduce(permutationsobject *po, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003295{
3296 if (po->result == NULL) {
3297 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3298 } else if (po->stopped) {
3299 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3300 } else {
3301 PyObject *indices=NULL, *cycles=NULL;
3302 Py_ssize_t n, i;
3303
3304 /* we must pickle the indices and cycles and use them for setstate */
3305 n = PyTuple_GET_SIZE(po->pool);
3306 indices = PyTuple_New(n);
3307 if (indices == NULL)
3308 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003309 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003310 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3311 if (!index)
3312 goto err;
3313 PyTuple_SET_ITEM(indices, i, index);
3314 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003315
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003316 cycles = PyTuple_New(po->r);
3317 if (cycles == NULL)
3318 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003319 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003320 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3321 if (!index)
3322 goto err;
3323 PyTuple_SET_ITEM(cycles, i, index);
3324 }
3325 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3326 po->pool, po->r,
3327 indices, cycles);
3328 err:
3329 Py_XDECREF(indices);
3330 Py_XDECREF(cycles);
3331 return NULL;
3332 }
3333}
3334
3335static PyObject *
3336permutations_setstate(permutationsobject *po, PyObject *state)
3337{
3338 PyObject *indices, *cycles, *result;
3339 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003340
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003341 if (!PyTuple_Check(state)) {
3342 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3343 return NULL;
3344 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003345 if (!PyArg_ParseTuple(state, "O!O!",
3346 &PyTuple_Type, &indices,
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003347 &PyTuple_Type, &cycles)) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003348 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003349 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003350
3351 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003352 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003353 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3354 return NULL;
3355 }
3356
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003357 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003358 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3359 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3360 if (index < 0 && PyErr_Occurred())
3361 return NULL; /* not an integer */
3362 /* clamp the index */
3363 if (index < 0)
3364 index = 0;
3365 else if (index > n-1)
3366 index = n-1;
3367 po->indices[i] = index;
3368 }
3369
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003370 for (i=0; i<po->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003371 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3372 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3373 if (index < 0 && PyErr_Occurred())
3374 return NULL; /* not an integer */
3375 if (index < 1)
3376 index = 1;
3377 else if (index > n-i)
3378 index = n-i;
3379 po->cycles[i] = index;
3380 }
3381 result = PyTuple_New(po->r);
3382 if (result == NULL)
3383 return NULL;
3384 for (i=0; i<po->r; i++) {
3385 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3386 Py_INCREF(element);
3387 PyTuple_SET_ITEM(result, i, element);
3388 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003389 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003390 Py_RETURN_NONE;
3391}
3392
3393static PyMethodDef permuations_methods[] = {
3394 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3395 reduce_doc},
3396 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3397 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003398 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3399 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003400 {NULL, NULL} /* sentinel */
3401};
3402
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003403static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003404 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003405 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003406 sizeof(permutationsobject), /* tp_basicsize */
3407 0, /* tp_itemsize */
3408 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003409 (destructor)permutations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003410 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003411 0, /* tp_getattr */
3412 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003413 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003414 0, /* tp_repr */
3415 0, /* tp_as_number */
3416 0, /* tp_as_sequence */
3417 0, /* tp_as_mapping */
3418 0, /* tp_hash */
3419 0, /* tp_call */
3420 0, /* tp_str */
3421 PyObject_GenericGetAttr, /* tp_getattro */
3422 0, /* tp_setattro */
3423 0, /* tp_as_buffer */
3424 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3425 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003426 itertools_permutations__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003427 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003428 0, /* tp_clear */
3429 0, /* tp_richcompare */
3430 0, /* tp_weaklistoffset */
3431 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003432 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003433 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003434 0, /* tp_members */
3435 0, /* tp_getset */
3436 0, /* tp_base */
3437 0, /* tp_dict */
3438 0, /* tp_descr_get */
3439 0, /* tp_descr_set */
3440 0, /* tp_dictoffset */
3441 0, /* tp_init */
3442 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003443 itertools_permutations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003444 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003445};
3446
Tal Einatc4bccd32018-09-12 00:49:13 +03003447
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003448/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003449
3450typedef struct {
3451 PyObject_HEAD
3452 PyObject *total;
3453 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003454 PyObject *binop;
Lisa Roach9718b592018-09-23 17:34:59 -07003455 PyObject *initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003456} accumulateobject;
3457
Tal Einatc4bccd32018-09-12 00:49:13 +03003458/*[clinic input]
3459@classmethod
3460itertools.accumulate.__new__
3461 iterable: object
3462 func as binop: object = None
Lisa Roach9718b592018-09-23 17:34:59 -07003463 *
3464 initial: object = None
Tal Einatc4bccd32018-09-12 00:49:13 +03003465Return series of accumulated sums (or other binary function results).
3466[clinic start generated code]*/
3467
Raymond Hettinger482ba772010-12-01 22:48:00 +00003468static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003469itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
Lisa Roach9718b592018-09-23 17:34:59 -07003470 PyObject *binop, PyObject *initial)
3471/*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003472{
Raymond Hettinger482ba772010-12-01 22:48:00 +00003473 PyObject *it;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003474 accumulateobject *lz;
3475
Raymond Hettinger482ba772010-12-01 22:48:00 +00003476 /* Get iterator. */
3477 it = PyObject_GetIter(iterable);
3478 if (it == NULL)
3479 return NULL;
3480
Raymond Hettinger482ba772010-12-01 22:48:00 +00003481 /* create accumulateobject structure */
3482 lz = (accumulateobject *)type->tp_alloc(type, 0);
3483 if (lz == NULL) {
3484 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003485 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003486 }
3487
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003488 if (binop != Py_None) {
3489 Py_XINCREF(binop);
3490 lz->binop = binop;
3491 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003492 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003493 lz->it = it;
Lisa Roach9718b592018-09-23 17:34:59 -07003494 Py_XINCREF(initial);
3495 lz->initial = initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003496 return (PyObject *)lz;
3497}
3498
3499static void
3500accumulate_dealloc(accumulateobject *lz)
3501{
3502 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003503 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003504 Py_XDECREF(lz->total);
3505 Py_XDECREF(lz->it);
Lisa Roach9718b592018-09-23 17:34:59 -07003506 Py_XDECREF(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003507 Py_TYPE(lz)->tp_free(lz);
3508}
3509
3510static int
3511accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3512{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003513 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003514 Py_VISIT(lz->it);
3515 Py_VISIT(lz->total);
Lisa Roach9718b592018-09-23 17:34:59 -07003516 Py_VISIT(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003517 return 0;
3518}
3519
3520static PyObject *
3521accumulate_next(accumulateobject *lz)
3522{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003523 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003524
Lisa Roach9718b592018-09-23 17:34:59 -07003525 if (lz->initial != Py_None) {
3526 lz->total = lz->initial;
3527 Py_INCREF(Py_None);
3528 lz->initial = Py_None;
3529 Py_INCREF(lz->total);
3530 return lz->total;
3531 }
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003532 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003533 if (val == NULL)
3534 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003535
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003536 if (lz->total == NULL) {
3537 Py_INCREF(val);
3538 lz->total = val;
3539 return lz->total;
3540 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003541
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003542 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003543 newtotal = PyNumber_Add(lz->total, val);
3544 else
3545 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003546 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003547 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003548 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003549
Raymond Hettinger482ba772010-12-01 22:48:00 +00003550 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003551 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003552 return newtotal;
3553}
3554
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003555static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303556accumulate_reduce(accumulateobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003557{
Lisa Roach9718b592018-09-23 17:34:59 -07003558 if (lz->initial != Py_None) {
3559 PyObject *it;
3560
3561 assert(lz->total == NULL);
3562 if (PyType_Ready(&chain_type) < 0)
3563 return NULL;
3564 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3565 lz->initial, lz->it);
3566 if (it == NULL)
3567 return NULL;
3568 return Py_BuildValue("O(NO)O", Py_TYPE(lz),
3569 it, lz->binop?lz->binop:Py_None, Py_None);
3570 }
Serhiy Storchakad5516252016-03-06 14:00:45 +02003571 if (lz->total == Py_None) {
3572 PyObject *it;
3573
3574 if (PyType_Ready(&chain_type) < 0)
3575 return NULL;
3576 if (PyType_Ready(&islice_type) < 0)
3577 return NULL;
3578 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3579 lz->total, lz->it);
3580 if (it == NULL)
3581 return NULL;
3582 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3583 it, lz->binop ? lz->binop : Py_None);
3584 if (it == NULL)
3585 return NULL;
3586 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3587 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003588 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3589 lz->it, lz->binop?lz->binop:Py_None,
3590 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003591}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003592
3593static PyObject *
3594accumulate_setstate(accumulateobject *lz, PyObject *state)
3595{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003596 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003597 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003598 Py_RETURN_NONE;
3599}
3600
3601static PyMethodDef accumulate_methods[] = {
3602 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3603 reduce_doc},
3604 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3605 setstate_doc},
3606 {NULL, NULL} /* sentinel */
3607};
3608
Raymond Hettinger482ba772010-12-01 22:48:00 +00003609static PyTypeObject accumulate_type = {
3610 PyVarObject_HEAD_INIT(NULL, 0)
3611 "itertools.accumulate", /* tp_name */
3612 sizeof(accumulateobject), /* tp_basicsize */
3613 0, /* tp_itemsize */
3614 /* methods */
3615 (destructor)accumulate_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003616 0, /* tp_vectorcall_offset */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003617 0, /* tp_getattr */
3618 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003619 0, /* tp_as_async */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003620 0, /* tp_repr */
3621 0, /* tp_as_number */
3622 0, /* tp_as_sequence */
3623 0, /* tp_as_mapping */
3624 0, /* tp_hash */
3625 0, /* tp_call */
3626 0, /* tp_str */
3627 PyObject_GenericGetAttr, /* tp_getattro */
3628 0, /* tp_setattro */
3629 0, /* tp_as_buffer */
3630 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3631 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003632 itertools_accumulate__doc__, /* tp_doc */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003633 (traverseproc)accumulate_traverse, /* tp_traverse */
3634 0, /* tp_clear */
3635 0, /* tp_richcompare */
3636 0, /* tp_weaklistoffset */
3637 PyObject_SelfIter, /* tp_iter */
3638 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003639 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003640 0, /* tp_members */
3641 0, /* tp_getset */
3642 0, /* tp_base */
3643 0, /* tp_dict */
3644 0, /* tp_descr_get */
3645 0, /* tp_descr_set */
3646 0, /* tp_dictoffset */
3647 0, /* tp_init */
3648 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003649 itertools_accumulate, /* tp_new */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003650 PyObject_GC_Del, /* tp_free */
3651};
3652
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003653
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003654/* compress object ************************************************************/
3655
3656/* Equivalent to:
3657
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003658 def compress(data, selectors):
3659 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3660 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003661*/
3662
3663typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003664 PyObject_HEAD
3665 PyObject *data;
3666 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003667} compressobject;
3668
Tal Einatc4bccd32018-09-12 00:49:13 +03003669/*[clinic input]
3670@classmethod
3671itertools.compress.__new__
3672 data as seq1: object
3673 selectors as seq2: object
3674Return data elements corresponding to true selector elements.
3675
3676Forms a shorter iterator from selected data elements using the selectors to
3677choose the data elements.
3678[clinic start generated code]*/
3679
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003680static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003681itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
3682/*[clinic end generated code: output=7e67157212ed09e0 input=79596d7cd20c77e5]*/
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003683{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003684 PyObject *data=NULL, *selectors=NULL;
3685 compressobject *lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003686
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003687 data = PyObject_GetIter(seq1);
3688 if (data == NULL)
3689 goto fail;
3690 selectors = PyObject_GetIter(seq2);
3691 if (selectors == NULL)
3692 goto fail;
3693
3694 /* create compressobject structure */
3695 lz = (compressobject *)type->tp_alloc(type, 0);
3696 if (lz == NULL)
3697 goto fail;
3698 lz->data = data;
3699 lz->selectors = selectors;
3700 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003701
3702fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003703 Py_XDECREF(data);
3704 Py_XDECREF(selectors);
3705 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003706}
3707
3708static void
3709compress_dealloc(compressobject *lz)
3710{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003711 PyObject_GC_UnTrack(lz);
3712 Py_XDECREF(lz->data);
3713 Py_XDECREF(lz->selectors);
3714 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003715}
3716
3717static int
3718compress_traverse(compressobject *lz, visitproc visit, void *arg)
3719{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003720 Py_VISIT(lz->data);
3721 Py_VISIT(lz->selectors);
3722 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003723}
3724
3725static PyObject *
3726compress_next(compressobject *lz)
3727{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003728 PyObject *data = lz->data, *selectors = lz->selectors;
3729 PyObject *datum, *selector;
3730 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3731 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3732 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003733
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003734 while (1) {
3735 /* Steps: get datum, get selector, evaluate selector.
3736 Order is important (to match the pure python version
3737 in terms of which input gets a chance to raise an
3738 exception first).
3739 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003740
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003741 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003742 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003743 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003744
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003745 selector = selectornext(selectors);
3746 if (selector == NULL) {
3747 Py_DECREF(datum);
3748 return NULL;
3749 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003750
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003751 ok = PyObject_IsTrue(selector);
3752 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003753 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003754 return datum;
3755 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003756 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003757 return NULL;
3758 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003759}
3760
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003761static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303762compress_reduce(compressobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003763{
3764 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3765 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003766}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003767
3768static PyMethodDef compress_methods[] = {
3769 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3770 reduce_doc},
3771 {NULL, NULL} /* sentinel */
3772};
3773
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003774static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003775 PyVarObject_HEAD_INIT(NULL, 0)
3776 "itertools.compress", /* tp_name */
3777 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003778 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003779 /* methods */
3780 (destructor)compress_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003781 0, /* tp_vectorcall_offset */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003782 0, /* tp_getattr */
3783 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003784 0, /* tp_as_async */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003785 0, /* tp_repr */
3786 0, /* tp_as_number */
3787 0, /* tp_as_sequence */
3788 0, /* tp_as_mapping */
3789 0, /* tp_hash */
3790 0, /* tp_call */
3791 0, /* tp_str */
3792 PyObject_GenericGetAttr, /* tp_getattro */
3793 0, /* tp_setattro */
3794 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003795 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003796 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003797 itertools_compress__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003798 (traverseproc)compress_traverse, /* tp_traverse */
3799 0, /* tp_clear */
3800 0, /* tp_richcompare */
3801 0, /* tp_weaklistoffset */
3802 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003803 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003804 compress_methods, /* tp_methods */
3805 0, /* tp_members */
3806 0, /* tp_getset */
3807 0, /* tp_base */
3808 0, /* tp_dict */
3809 0, /* tp_descr_get */
3810 0, /* tp_descr_set */
3811 0, /* tp_dictoffset */
3812 0, /* tp_init */
3813 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003814 itertools_compress, /* tp_new */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003815 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003816};
3817
3818
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003819/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003820
3821typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003822 PyObject_HEAD
3823 PyObject *func;
3824 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003825} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003826
Tal Einatc4bccd32018-09-12 00:49:13 +03003827/*[clinic input]
3828@classmethod
3829itertools.filterfalse.__new__
3830 function as func: object
3831 iterable as seq: object
3832 /
3833Return those items of iterable for which function(item) is false.
3834
3835If function is None, return the items that are false.
3836[clinic start generated code]*/
3837
Raymond Hettinger60eca932003-02-09 06:40:58 +00003838static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003839itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
3840/*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003841{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003842 PyObject *it;
3843 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003844
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003845 /* Get iterator. */
3846 it = PyObject_GetIter(seq);
3847 if (it == NULL)
3848 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003849
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003850 /* create filterfalseobject structure */
3851 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3852 if (lz == NULL) {
3853 Py_DECREF(it);
3854 return NULL;
3855 }
3856 Py_INCREF(func);
3857 lz->func = func;
3858 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003859
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003860 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003861}
3862
3863static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003864filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003865{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003866 PyObject_GC_UnTrack(lz);
3867 Py_XDECREF(lz->func);
3868 Py_XDECREF(lz->it);
3869 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003870}
3871
3872static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003873filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003874{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003875 Py_VISIT(lz->it);
3876 Py_VISIT(lz->func);
3877 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003878}
3879
3880static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003881filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003882{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003883 PyObject *item;
3884 PyObject *it = lz->it;
3885 long ok;
3886 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003887
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003888 iternext = *Py_TYPE(it)->tp_iternext;
3889 for (;;) {
3890 item = iternext(it);
3891 if (item == NULL)
3892 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003893
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003894 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3895 ok = PyObject_IsTrue(item);
3896 } else {
3897 PyObject *good;
Petr Viktorinffd97532020-02-11 17:46:57 +01003898 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003899 if (good == NULL) {
3900 Py_DECREF(item);
3901 return NULL;
3902 }
3903 ok = PyObject_IsTrue(good);
3904 Py_DECREF(good);
3905 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003906 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003907 return item;
3908 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003909 if (ok < 0)
3910 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003911 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003912}
3913
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003914static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303915filterfalse_reduce(filterfalseobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003916{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003917 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3918}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003919
3920static PyMethodDef filterfalse_methods[] = {
3921 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3922 reduce_doc},
3923 {NULL, NULL} /* sentinel */
3924};
3925
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003926static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003927 PyVarObject_HEAD_INIT(NULL, 0)
3928 "itertools.filterfalse", /* tp_name */
3929 sizeof(filterfalseobject), /* tp_basicsize */
3930 0, /* tp_itemsize */
3931 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003932 (destructor)filterfalse_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003933 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003934 0, /* tp_getattr */
3935 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003936 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003937 0, /* tp_repr */
3938 0, /* tp_as_number */
3939 0, /* tp_as_sequence */
3940 0, /* tp_as_mapping */
3941 0, /* tp_hash */
3942 0, /* tp_call */
3943 0, /* tp_str */
3944 PyObject_GenericGetAttr, /* tp_getattro */
3945 0, /* tp_setattro */
3946 0, /* tp_as_buffer */
3947 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3948 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003949 itertools_filterfalse__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003950 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003951 0, /* tp_clear */
3952 0, /* tp_richcompare */
3953 0, /* tp_weaklistoffset */
3954 PyObject_SelfIter, /* tp_iter */
3955 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003956 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003957 0, /* tp_members */
3958 0, /* tp_getset */
3959 0, /* tp_base */
3960 0, /* tp_dict */
3961 0, /* tp_descr_get */
3962 0, /* tp_descr_set */
3963 0, /* tp_dictoffset */
3964 0, /* tp_init */
3965 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003966 itertools_filterfalse, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003967 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003968};
3969
3970
3971/* count object ************************************************************/
3972
3973typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003974 PyObject_HEAD
3975 Py_ssize_t cnt;
3976 PyObject *long_cnt;
3977 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003978} countobject;
3979
Raymond Hettinger30729212009-02-12 06:28:27 +00003980/* Counting logic and invariants:
3981
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003982fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003983
Serhiy Storchaka483405b2015-02-17 10:14:30 +02003984 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003985 Advances with: cnt += 1
3986 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003987
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003988slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003989
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003990 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3991 All counting is done with python objects (no overflows or underflows).
3992 Advances with: long_cnt += long_step
3993 Step may be zero -- effectively a slow version of repeat(cnt).
3994 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003995*/
3996
Tal Einatc4bccd32018-09-12 00:49:13 +03003997/*[clinic input]
3998@classmethod
3999itertools.count.__new__
4000 start as long_cnt: object(c_default="NULL") = 0
4001 step as long_step: object(c_default="NULL") = 1
4002Return a count object whose .__next__() method returns consecutive values.
4003
4004Equivalent to:
4005 def count(firstval=0, step=1):
4006 x = firstval
4007 while 1:
4008 yield x
4009 x += step
4010[clinic start generated code]*/
4011
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004012static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03004013itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
4014 PyObject *long_step)
4015/*[clinic end generated code: output=09a9250aebd00b1c input=d7a85eec18bfcd94]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004016{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004017 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004018 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004019 Py_ssize_t cnt = 0;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00004020 long step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004021
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004022 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
4023 (long_step != NULL && !PyNumber_Check(long_step))) {
4024 PyErr_SetString(PyExc_TypeError, "a number is required");
4025 return NULL;
4026 }
Raymond Hettinger30729212009-02-12 06:28:27 +00004027
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004028 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
4029 (long_step == NULL || PyLong_Check(long_step));
4030
4031 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004032 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004033 if (fast_mode) {
4034 assert(PyLong_Check(long_cnt));
4035 cnt = PyLong_AsSsize_t(long_cnt);
4036 if (cnt == -1 && PyErr_Occurred()) {
4037 PyErr_Clear();
4038 fast_mode = 0;
4039 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004040 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004041 } else {
4042 cnt = 0;
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004043 long_cnt = _PyLong_Zero;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004044 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004045 Py_INCREF(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004046
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004047 /* If not specified, step defaults to 1 */
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004048 if (long_step == NULL)
4049 long_step = _PyLong_One;
4050 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004051
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004052 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004053
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004054 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004055 if (fast_mode) {
4056 assert(PyLong_Check(long_step));
4057 step = PyLong_AsLong(long_step);
4058 if (step != 1) {
4059 fast_mode = 0;
4060 if (step == -1 && PyErr_Occurred())
4061 PyErr_Clear();
4062 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004063 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004064
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004065 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004066 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004067 else
4068 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004069
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004070 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4071 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4072 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004073 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004074
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004075 /* create countobject structure */
4076 lz = (countobject *)type->tp_alloc(type, 0);
4077 if (lz == NULL) {
4078 Py_XDECREF(long_cnt);
Zackery Spytz0523c392019-03-26 00:05:29 -06004079 Py_DECREF(long_step);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004080 return NULL;
4081 }
4082 lz->cnt = cnt;
4083 lz->long_cnt = long_cnt;
4084 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004085
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004086 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004087}
4088
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004089static void
4090count_dealloc(countobject *lz)
4091{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004092 PyObject_GC_UnTrack(lz);
4093 Py_XDECREF(lz->long_cnt);
4094 Py_XDECREF(lz->long_step);
4095 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004096}
4097
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004098static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004099count_traverse(countobject *lz, visitproc visit, void *arg)
4100{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004101 Py_VISIT(lz->long_cnt);
4102 Py_VISIT(lz->long_step);
4103 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004104}
4105
4106static PyObject *
4107count_nextlong(countobject *lz)
4108{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004109 PyObject *long_cnt;
4110 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004111
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004112 long_cnt = lz->long_cnt;
4113 if (long_cnt == NULL) {
4114 /* Switch to slow_mode */
4115 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4116 if (long_cnt == NULL)
4117 return NULL;
4118 }
4119 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004120
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004121 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4122 if (stepped_up == NULL)
4123 return NULL;
4124 lz->long_cnt = stepped_up;
4125 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004126}
4127
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004128static PyObject *
4129count_next(countobject *lz)
4130{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004131 if (lz->cnt == PY_SSIZE_T_MAX)
4132 return count_nextlong(lz);
4133 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004134}
4135
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004136static PyObject *
4137count_repr(countobject *lz)
4138{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004139 if (lz->cnt != PY_SSIZE_T_MAX)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004140 return PyUnicode_FromFormat("%s(%zd)",
4141 _PyType_Name(Py_TYPE(lz)), lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004142
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004143 if (PyLong_Check(lz->long_step)) {
4144 long step = PyLong_AsLong(lz->long_step);
4145 if (step == -1 && PyErr_Occurred()) {
4146 PyErr_Clear();
4147 }
4148 if (step == 1) {
4149 /* Don't display step when it is an integer equal to 1 */
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004150 return PyUnicode_FromFormat("%s(%R)",
4151 _PyType_Name(Py_TYPE(lz)),
4152 lz->long_cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004153 }
4154 }
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004155 return PyUnicode_FromFormat("%s(%R, %R)",
4156 _PyType_Name(Py_TYPE(lz)),
4157 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004158}
4159
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004160static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304161count_reduce(countobject *lz, PyObject *Py_UNUSED(ignored))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004162{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004163 if (lz->cnt == PY_SSIZE_T_MAX)
4164 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4165 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004166}
4167
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004168static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004169 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004170 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004171 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004172};
4173
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004174static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004175 PyVarObject_HEAD_INIT(NULL, 0)
4176 "itertools.count", /* tp_name */
4177 sizeof(countobject), /* tp_basicsize */
4178 0, /* tp_itemsize */
4179 /* methods */
4180 (destructor)count_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004181 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004182 0, /* tp_getattr */
4183 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004184 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004185 (reprfunc)count_repr, /* tp_repr */
4186 0, /* tp_as_number */
4187 0, /* tp_as_sequence */
4188 0, /* tp_as_mapping */
4189 0, /* tp_hash */
4190 0, /* tp_call */
4191 0, /* tp_str */
4192 PyObject_GenericGetAttr, /* tp_getattro */
4193 0, /* tp_setattro */
4194 0, /* tp_as_buffer */
4195 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004196 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03004197 itertools_count__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004198 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004199 0, /* tp_clear */
4200 0, /* tp_richcompare */
4201 0, /* tp_weaklistoffset */
4202 PyObject_SelfIter, /* tp_iter */
4203 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004204 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004205 0, /* tp_members */
4206 0, /* tp_getset */
4207 0, /* tp_base */
4208 0, /* tp_dict */
4209 0, /* tp_descr_get */
4210 0, /* tp_descr_set */
4211 0, /* tp_dictoffset */
4212 0, /* tp_init */
4213 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03004214 itertools_count, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004215 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004216};
4217
4218
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004219/* repeat object ************************************************************/
4220
4221typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004222 PyObject_HEAD
4223 PyObject *element;
4224 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004225} repeatobject;
4226
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004227static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004228
4229static PyObject *
4230repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4231{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004232 repeatobject *ro;
4233 PyObject *element;
Raymond Hettingeraf464502019-11-09 20:28:31 -08004234 Py_ssize_t cnt = -1, n_args;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004235 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004236
Raymond Hettingeraf464502019-11-09 20:28:31 -08004237 n_args = PyTuple_GET_SIZE(args);
4238 if (kwds != NULL)
4239 n_args += PyDict_GET_SIZE(kwds);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004240 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4241 &element, &cnt))
4242 return NULL;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004243 /* Does user supply times argument? */
Raymond Hettingeraf464502019-11-09 20:28:31 -08004244 if (n_args == 2 && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004245 cnt = 0;
4246
4247 ro = (repeatobject *)type->tp_alloc(type, 0);
4248 if (ro == NULL)
4249 return NULL;
4250 Py_INCREF(element);
4251 ro->element = element;
4252 ro->cnt = cnt;
4253 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004254}
4255
4256static void
4257repeat_dealloc(repeatobject *ro)
4258{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004259 PyObject_GC_UnTrack(ro);
4260 Py_XDECREF(ro->element);
4261 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004262}
4263
4264static int
4265repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4266{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004267 Py_VISIT(ro->element);
4268 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004269}
4270
4271static PyObject *
4272repeat_next(repeatobject *ro)
4273{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004274 if (ro->cnt == 0)
4275 return NULL;
4276 if (ro->cnt > 0)
4277 ro->cnt--;
4278 Py_INCREF(ro->element);
4279 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004280}
4281
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004282static PyObject *
4283repeat_repr(repeatobject *ro)
4284{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004285 if (ro->cnt == -1)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004286 return PyUnicode_FromFormat("%s(%R)",
4287 _PyType_Name(Py_TYPE(ro)), ro->element);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004288 else
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004289 return PyUnicode_FromFormat("%s(%R, %zd)",
4290 _PyType_Name(Py_TYPE(ro)), ro->element,
4291 ro->cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004292}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004293
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004294static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304295repeat_len(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004296{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004297 if (ro->cnt == -1) {
4298 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4299 return NULL;
4300 }
4301 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004302}
4303
Armin Rigof5b3e362006-02-11 21:32:43 +00004304PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004305
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004306static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304307repeat_reduce(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004308{
4309 /* unpickle this so that a new repeat iterator is constructed with an
4310 * object, then call __setstate__ on it to set cnt
4311 */
4312 if (ro->cnt >= 0)
4313 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4314 else
4315 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4316}
4317
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004318static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004319 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004320 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004321 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004322};
4323
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004324PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004325"repeat(object [,times]) -> create an iterator which returns the object\n\
4326for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004327endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004328
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004329static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004330 PyVarObject_HEAD_INIT(NULL, 0)
4331 "itertools.repeat", /* tp_name */
4332 sizeof(repeatobject), /* tp_basicsize */
4333 0, /* tp_itemsize */
4334 /* methods */
4335 (destructor)repeat_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004336 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004337 0, /* tp_getattr */
4338 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004339 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004340 (reprfunc)repeat_repr, /* tp_repr */
4341 0, /* tp_as_number */
4342 0, /* tp_as_sequence */
4343 0, /* tp_as_mapping */
4344 0, /* tp_hash */
4345 0, /* tp_call */
4346 0, /* tp_str */
4347 PyObject_GenericGetAttr, /* tp_getattro */
4348 0, /* tp_setattro */
4349 0, /* tp_as_buffer */
4350 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4351 Py_TPFLAGS_BASETYPE, /* tp_flags */
4352 repeat_doc, /* tp_doc */
4353 (traverseproc)repeat_traverse, /* tp_traverse */
4354 0, /* tp_clear */
4355 0, /* tp_richcompare */
4356 0, /* tp_weaklistoffset */
4357 PyObject_SelfIter, /* tp_iter */
4358 (iternextfunc)repeat_next, /* tp_iternext */
4359 repeat_methods, /* tp_methods */
4360 0, /* tp_members */
4361 0, /* tp_getset */
4362 0, /* tp_base */
4363 0, /* tp_dict */
4364 0, /* tp_descr_get */
4365 0, /* tp_descr_set */
4366 0, /* tp_dictoffset */
4367 0, /* tp_init */
4368 0, /* tp_alloc */
4369 repeat_new, /* tp_new */
4370 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004371};
4372
Tal Einatc4bccd32018-09-12 00:49:13 +03004373
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004374/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004375
4376typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004377 PyObject_HEAD
4378 Py_ssize_t tuplesize;
4379 Py_ssize_t numactive;
4380 PyObject *ittuple; /* tuple of iterators */
4381 PyObject *result;
4382 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004383} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004384
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004385static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004386
4387static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004388zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004389{
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004390 _Py_IDENTIFIER(fillvalue);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004391 ziplongestobject *lz;
4392 Py_ssize_t i;
4393 PyObject *ittuple; /* tuple of iterators */
4394 PyObject *result;
4395 PyObject *fillvalue = Py_None;
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004396 Py_ssize_t tuplesize;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004397
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004398 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004399 fillvalue = NULL;
4400 if (PyDict_GET_SIZE(kwds) == 1) {
4401 fillvalue = _PyDict_GetItemIdWithError(kwds, &PyId_fillvalue);
4402 }
4403 if (fillvalue == NULL) {
4404 if (!PyErr_Occurred()) {
4405 PyErr_SetString(PyExc_TypeError,
4406 "zip_longest() got an unexpected keyword argument");
4407 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004408 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004409 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004410 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004411
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004412 /* args must be a tuple */
4413 assert(PyTuple_Check(args));
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004414 tuplesize = PyTuple_GET_SIZE(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004415
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004416 /* obtain iterators */
4417 ittuple = PyTuple_New(tuplesize);
4418 if (ittuple == NULL)
4419 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004420 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004421 PyObject *item = PyTuple_GET_ITEM(args, i);
4422 PyObject *it = PyObject_GetIter(item);
4423 if (it == NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004424 Py_DECREF(ittuple);
4425 return NULL;
4426 }
4427 PyTuple_SET_ITEM(ittuple, i, it);
4428 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004429
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004430 /* create a result holder */
4431 result = PyTuple_New(tuplesize);
4432 if (result == NULL) {
4433 Py_DECREF(ittuple);
4434 return NULL;
4435 }
4436 for (i=0 ; i < tuplesize ; i++) {
4437 Py_INCREF(Py_None);
4438 PyTuple_SET_ITEM(result, i, Py_None);
4439 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004440
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004441 /* create ziplongestobject structure */
4442 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4443 if (lz == NULL) {
4444 Py_DECREF(ittuple);
4445 Py_DECREF(result);
4446 return NULL;
4447 }
4448 lz->ittuple = ittuple;
4449 lz->tuplesize = tuplesize;
4450 lz->numactive = tuplesize;
4451 lz->result = result;
4452 Py_INCREF(fillvalue);
4453 lz->fillvalue = fillvalue;
4454 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004455}
4456
4457static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004458zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004459{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004460 PyObject_GC_UnTrack(lz);
4461 Py_XDECREF(lz->ittuple);
4462 Py_XDECREF(lz->result);
4463 Py_XDECREF(lz->fillvalue);
4464 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004465}
4466
4467static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004468zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004469{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004470 Py_VISIT(lz->ittuple);
4471 Py_VISIT(lz->result);
4472 Py_VISIT(lz->fillvalue);
4473 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004474}
4475
4476static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004477zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004478{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004479 Py_ssize_t i;
4480 Py_ssize_t tuplesize = lz->tuplesize;
4481 PyObject *result = lz->result;
4482 PyObject *it;
4483 PyObject *item;
4484 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004485
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004486 if (tuplesize == 0)
4487 return NULL;
4488 if (lz->numactive == 0)
4489 return NULL;
4490 if (Py_REFCNT(result) == 1) {
4491 Py_INCREF(result);
4492 for (i=0 ; i < tuplesize ; i++) {
4493 it = PyTuple_GET_ITEM(lz->ittuple, i);
4494 if (it == NULL) {
4495 Py_INCREF(lz->fillvalue);
4496 item = lz->fillvalue;
4497 } else {
4498 item = PyIter_Next(it);
4499 if (item == NULL) {
4500 lz->numactive -= 1;
4501 if (lz->numactive == 0 || PyErr_Occurred()) {
4502 lz->numactive = 0;
4503 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004504 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004505 } else {
4506 Py_INCREF(lz->fillvalue);
4507 item = lz->fillvalue;
4508 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4509 Py_DECREF(it);
4510 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004511 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004512 }
4513 olditem = PyTuple_GET_ITEM(result, i);
4514 PyTuple_SET_ITEM(result, i, item);
4515 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004516 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004517 } else {
4518 result = PyTuple_New(tuplesize);
4519 if (result == NULL)
4520 return NULL;
4521 for (i=0 ; i < tuplesize ; i++) {
4522 it = PyTuple_GET_ITEM(lz->ittuple, i);
4523 if (it == NULL) {
4524 Py_INCREF(lz->fillvalue);
4525 item = lz->fillvalue;
4526 } else {
4527 item = PyIter_Next(it);
4528 if (item == NULL) {
4529 lz->numactive -= 1;
4530 if (lz->numactive == 0 || PyErr_Occurred()) {
4531 lz->numactive = 0;
4532 Py_DECREF(result);
4533 return NULL;
4534 } else {
4535 Py_INCREF(lz->fillvalue);
4536 item = lz->fillvalue;
4537 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4538 Py_DECREF(it);
4539 }
4540 }
4541 }
4542 PyTuple_SET_ITEM(result, i, item);
4543 }
4544 }
4545 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004546}
4547
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004548static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304549zip_longest_reduce(ziplongestobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004550{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004551
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004552 /* Create a new tuple with empty sequences where appropriate to pickle.
4553 * Then use setstate to set the fillvalue
4554 */
4555 int i;
4556 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004557
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004558 if (args == NULL)
4559 return NULL;
4560 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4561 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4562 if (elem == NULL) {
4563 elem = PyTuple_New(0);
4564 if (elem == NULL) {
4565 Py_DECREF(args);
4566 return NULL;
4567 }
4568 } else
4569 Py_INCREF(elem);
4570 PyTuple_SET_ITEM(args, i, elem);
4571 }
4572 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4573}
4574
4575static PyObject *
4576zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4577{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004578 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004579 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004580 Py_RETURN_NONE;
4581}
4582
4583static PyMethodDef zip_longest_methods[] = {
4584 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4585 reduce_doc},
4586 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4587 setstate_doc},
4588 {NULL, NULL} /* sentinel */
4589};
4590
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004591PyDoc_STRVAR(zip_longest_doc,
4592"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004593\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004594Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004595the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004596method continues until the longest iterable in the argument sequence\n\
4597is exhausted and then it raises StopIteration. When the shorter iterables\n\
4598are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4599defaults to None or can be specified by a keyword argument.\n\
4600");
4601
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004602static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004603 PyVarObject_HEAD_INIT(NULL, 0)
4604 "itertools.zip_longest", /* tp_name */
4605 sizeof(ziplongestobject), /* tp_basicsize */
4606 0, /* tp_itemsize */
4607 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004608 (destructor)zip_longest_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004609 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004610 0, /* tp_getattr */
4611 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004612 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004613 0, /* tp_repr */
4614 0, /* tp_as_number */
4615 0, /* tp_as_sequence */
4616 0, /* tp_as_mapping */
4617 0, /* tp_hash */
4618 0, /* tp_call */
4619 0, /* tp_str */
4620 PyObject_GenericGetAttr, /* tp_getattro */
4621 0, /* tp_setattro */
4622 0, /* tp_as_buffer */
4623 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4624 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004625 zip_longest_doc, /* tp_doc */
4626 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004627 0, /* tp_clear */
4628 0, /* tp_richcompare */
4629 0, /* tp_weaklistoffset */
4630 PyObject_SelfIter, /* tp_iter */
4631 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004632 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004633 0, /* tp_members */
4634 0, /* tp_getset */
4635 0, /* tp_base */
4636 0, /* tp_dict */
4637 0, /* tp_descr_get */
4638 0, /* tp_descr_set */
4639 0, /* tp_dictoffset */
4640 0, /* tp_init */
4641 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004642 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004643 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004644};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004645
Tal Einatc4bccd32018-09-12 00:49:13 +03004646
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004647/* module level code ********************************************************/
4648
4649PyDoc_STRVAR(module_doc,
4650"Functional tools for creating and using iterators.\n\
4651\n\
4652Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004653count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004654cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004655repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004656\n\
4657Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004658accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
oldkaa0735f2018-02-02 16:52:55 +08004659chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
4660chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004661compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4662dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4663groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004664filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004665islice(seq, [start,] stop [, step]) --> elements from\n\
4666 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004667starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004668tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004669takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
oldkaa0735f2018-02-02 16:52:55 +08004670zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004671\n\
4672Combinatoric generators:\n\
4673product(p, q, ... [repeat=1]) --> cartesian product\n\
4674permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004675combinations(p, r)\n\
4676combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004677");
4678
Dong-hee Na514c4692020-03-18 02:46:24 +09004679static int
4680itertoolsmodule_exec(PyObject *m)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004681{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004682 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004683 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004684 &combinations_type,
4685 &cwr_type,
4686 &cycle_type,
4687 &dropwhile_type,
4688 &takewhile_type,
4689 &islice_type,
4690 &starmap_type,
4691 &chain_type,
4692 &compress_type,
4693 &filterfalse_type,
4694 &count_type,
4695 &ziplongest_type,
4696 &permutations_type,
4697 &product_type,
4698 &repeat_type,
4699 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004700 &_grouper_type,
4701 &tee_type,
Dong-hee Na05e4a292020-03-23 01:17:34 +09004702 &teedataobject_type
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004703 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004704
Victor Stinnerd2ec81a2020-02-07 09:17:07 +01004705 Py_SET_TYPE(&teedataobject_type, &PyType_Type);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004706
Dong-hee Na05e4a292020-03-23 01:17:34 +09004707 for (size_t i = 0; i < Py_ARRAY_LENGTH(typelist); i++) {
4708 if (PyModule_AddType(m, typelist[i]) < 0) {
Dong-hee Na514c4692020-03-18 02:46:24 +09004709 return -1;
4710 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004711 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004712
Dong-hee Na514c4692020-03-18 02:46:24 +09004713 return 0;
4714}
4715
4716static struct PyModuleDef_Slot itertoolsmodule_slots[] = {
4717 {Py_mod_exec, itertoolsmodule_exec},
4718 {0, NULL}
4719};
4720
4721static PyMethodDef module_methods[] = {
4722 ITERTOOLS_TEE_METHODDEF
4723 {NULL, NULL} /* sentinel */
4724};
4725
4726
4727static struct PyModuleDef itertoolsmodule = {
4728 PyModuleDef_HEAD_INIT,
4729 "itertools",
4730 module_doc,
4731 0,
4732 module_methods,
4733 itertoolsmodule_slots,
4734 NULL,
4735 NULL,
4736 NULL
4737};
4738
4739PyMODINIT_FUNC
4740PyInit_itertools(void)
4741{
4742 return PyModuleDef_Init(&itertoolsmodule);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004743}