blob: e60ad5bec43c6bd7a91115c1b54524f3780f4f64 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
Raymond Hettingerca3788c2015-08-16 14:49:24 -07002#define PY_SSIZE_T_CLEAN
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003#include "Python.h"
Sergey Fedoseev234531b2019-02-25 21:59:12 +05004#include "pycore_tupleobject.h"
Fred Drake08ebfec2004-10-17 19:36:57 +00005#include "structmember.h"
Raymond Hettinger96ef8112003-02-01 00:10:11 +00006
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00007/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +00008 by Raymond D. Hettinger <python@rcn.com>
Raymond Hettinger96ef8112003-02-01 00:10:11 +00009*/
10
Tal Einat3286ce42018-09-10 21:33:08 +030011/*[clinic input]
12module itertools
13class itertools.groupby "groupbyobject *" "&groupby_type"
14class itertools._grouper "_grouperobject *" "&_grouper_type"
Tal Einatc4bccd32018-09-12 00:49:13 +030015class itertools.teedataobject "teedataobject *" "&teedataobject_type"
16class itertools._tee "teeobject *" "&tee_type"
17class itertools.cycle "cycleobject *" "&cycle_type"
18class itertools.dropwhile "dropwhileobject *" "&dropwhile_type"
19class itertools.takewhile "takewhileobject *" "&takewhile_type"
20class itertools.starmap "starmapobject *" "&starmap_type"
21class itertools.chain "chainobject *" "&chain_type"
22class itertools.combinations "combinationsobject *" "&combinations_type"
23class itertools.combinations_with_replacement "cwr_object *" "&cwr_type"
24class itertools.permutations "permutationsobject *" "&permutations_type"
25class itertools.accumulate "accumulateobject *" "&accumulate_type"
26class itertools.compress "compressobject *" "&compress_type"
27class itertools.filterfalse "filterfalseobject *" "&filterfalse_type"
28class itertools.count "countobject *" "&count_type"
Tal Einat3286ce42018-09-10 21:33:08 +030029[clinic start generated code]*/
Tal Einatc4bccd32018-09-12 00:49:13 +030030/*[clinic end generated code: output=da39a3ee5e6b4b0d input=ea05c93c6d94726a]*/
Tal Einat3286ce42018-09-10 21:33:08 +030031
32static PyTypeObject groupby_type;
33static PyTypeObject _grouper_type;
Tal Einatc4bccd32018-09-12 00:49:13 +030034static PyTypeObject teedataobject_type;
35static PyTypeObject tee_type;
36static PyTypeObject cycle_type;
37static PyTypeObject dropwhile_type;
38static PyTypeObject takewhile_type;
39static PyTypeObject starmap_type;
40static PyTypeObject combinations_type;
41static PyTypeObject cwr_type;
42static PyTypeObject permutations_type;
43static PyTypeObject accumulate_type;
44static PyTypeObject compress_type;
45static PyTypeObject filterfalse_type;
46static PyTypeObject count_type;
47
Tal Einat3286ce42018-09-10 21:33:08 +030048#include "clinic/itertoolsmodule.c.h"
49
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000050
Raymond Hettingera6ea44a2015-08-17 23:55:28 -070051/* groupby object ************************************************************/
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000052
53typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000054 PyObject_HEAD
55 PyObject *it;
56 PyObject *keyfunc;
57 PyObject *tgtkey;
58 PyObject *currkey;
59 PyObject *currvalue;
Serhiy Storchakac247caf2017-09-24 13:36:11 +030060 const void *currgrouper; /* borrowed reference */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000061} groupbyobject;
62
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000063static PyObject *_grouper_create(groupbyobject *, PyObject *);
64
Tal Einat3286ce42018-09-10 21:33:08 +030065/*[clinic input]
66@classmethod
67itertools.groupby.__new__
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000068
Tal Einat3286ce42018-09-10 21:33:08 +030069 iterable as it: object
70 Elements to divide into groups according to the key function.
71 key as keyfunc: object = None
72 A function for computing the group category for each element.
73 If the key function is not specified or is None, the element itself
74 is used for grouping.
75
76make an iterator that returns consecutive keys and groups from the iterable
77[clinic start generated code]*/
78
79static PyObject *
80itertools_groupby_impl(PyTypeObject *type, PyObject *it, PyObject *keyfunc)
81/*[clinic end generated code: output=cbb1ae3a90fd4141 input=6b3d123e87ff65a1]*/
82{
83 groupbyobject *gbo;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000084
85 gbo = (groupbyobject *)type->tp_alloc(type, 0);
86 if (gbo == NULL)
87 return NULL;
88 gbo->tgtkey = NULL;
89 gbo->currkey = NULL;
90 gbo->currvalue = NULL;
91 gbo->keyfunc = keyfunc;
92 Py_INCREF(keyfunc);
93 gbo->it = PyObject_GetIter(it);
94 if (gbo->it == NULL) {
95 Py_DECREF(gbo);
96 return NULL;
97 }
98 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000099}
100
101static void
102groupby_dealloc(groupbyobject *gbo)
103{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000104 PyObject_GC_UnTrack(gbo);
105 Py_XDECREF(gbo->it);
106 Py_XDECREF(gbo->keyfunc);
107 Py_XDECREF(gbo->tgtkey);
108 Py_XDECREF(gbo->currkey);
109 Py_XDECREF(gbo->currvalue);
110 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000111}
112
113static int
114groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
115{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000116 Py_VISIT(gbo->it);
117 Py_VISIT(gbo->keyfunc);
118 Py_VISIT(gbo->tgtkey);
119 Py_VISIT(gbo->currkey);
120 Py_VISIT(gbo->currvalue);
121 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000122}
123
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300124Py_LOCAL_INLINE(int)
125groupby_step(groupbyobject *gbo)
126{
127 PyObject *newvalue, *newkey, *oldvalue;
128
129 newvalue = PyIter_Next(gbo->it);
130 if (newvalue == NULL)
131 return -1;
132
133 if (gbo->keyfunc == Py_None) {
134 newkey = newvalue;
135 Py_INCREF(newvalue);
136 } else {
Jeroen Demeyer196a5302019-07-04 12:31:34 +0200137 newkey = _PyObject_CallOneArg(gbo->keyfunc, newvalue);
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300138 if (newkey == NULL) {
139 Py_DECREF(newvalue);
140 return -1;
141 }
142 }
143
144 oldvalue = gbo->currvalue;
145 gbo->currvalue = newvalue;
146 Py_XSETREF(gbo->currkey, newkey);
147 Py_XDECREF(oldvalue);
148 return 0;
149}
150
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000151static PyObject *
152groupby_next(groupbyobject *gbo)
153{
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300154 PyObject *r, *grouper;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000155
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300156 gbo->currgrouper = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000157 /* skip to next iteration group */
158 for (;;) {
159 if (gbo->currkey == NULL)
160 /* pass */;
161 else if (gbo->tgtkey == NULL)
162 break;
163 else {
164 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000165
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700166 rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000167 if (rcmp == -1)
168 return NULL;
169 else if (rcmp == 0)
170 break;
171 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000172
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300173 if (groupby_step(gbo) < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000174 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000175 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000176 Py_INCREF(gbo->currkey);
Serhiy Storchakaec397562016-04-06 09:50:03 +0300177 Py_XSETREF(gbo->tgtkey, gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000178
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000179 grouper = _grouper_create(gbo, gbo->tgtkey);
180 if (grouper == NULL)
181 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000182
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000183 r = PyTuple_Pack(2, gbo->currkey, grouper);
184 Py_DECREF(grouper);
185 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000186}
187
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000188static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530189groupby_reduce(groupbyobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000190{
191 /* reduce as a 'new' call with an optional 'setstate' if groupby
192 * has started
193 */
194 PyObject *value;
195 if (lz->tgtkey && lz->currkey && lz->currvalue)
196 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
197 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
198 else
199 value = Py_BuildValue("O(OO)", Py_TYPE(lz),
200 lz->it, lz->keyfunc);
201
202 return value;
203}
204
205PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
206
207static PyObject *
208groupby_setstate(groupbyobject *lz, PyObject *state)
209{
210 PyObject *currkey, *currvalue, *tgtkey;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300211 if (!PyTuple_Check(state)) {
212 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000213 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300214 }
215 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey)) {
216 return NULL;
217 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200218 Py_INCREF(currkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300219 Py_XSETREF(lz->currkey, currkey);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200220 Py_INCREF(currvalue);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300221 Py_XSETREF(lz->currvalue, currvalue);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200222 Py_INCREF(tgtkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300223 Py_XSETREF(lz->tgtkey, tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000224 Py_RETURN_NONE;
225}
226
227PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
228
229static PyMethodDef groupby_methods[] = {
230 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
231 reduce_doc},
232 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
233 setstate_doc},
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700234 {NULL, NULL} /* sentinel */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000235};
236
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000237static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000238 PyVarObject_HEAD_INIT(NULL, 0)
239 "itertools.groupby", /* tp_name */
240 sizeof(groupbyobject), /* tp_basicsize */
241 0, /* tp_itemsize */
242 /* methods */
243 (destructor)groupby_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200244 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000245 0, /* tp_getattr */
246 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200247 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000248 0, /* tp_repr */
249 0, /* tp_as_number */
250 0, /* tp_as_sequence */
251 0, /* tp_as_mapping */
252 0, /* tp_hash */
253 0, /* tp_call */
254 0, /* tp_str */
255 PyObject_GenericGetAttr, /* tp_getattro */
256 0, /* tp_setattro */
257 0, /* tp_as_buffer */
258 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
259 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einat3286ce42018-09-10 21:33:08 +0300260 itertools_groupby__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000261 (traverseproc)groupby_traverse, /* tp_traverse */
262 0, /* tp_clear */
263 0, /* tp_richcompare */
264 0, /* tp_weaklistoffset */
265 PyObject_SelfIter, /* tp_iter */
266 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000267 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000268 0, /* tp_members */
269 0, /* tp_getset */
270 0, /* tp_base */
271 0, /* tp_dict */
272 0, /* tp_descr_get */
273 0, /* tp_descr_set */
274 0, /* tp_dictoffset */
275 0, /* tp_init */
276 0, /* tp_alloc */
Tal Einat3286ce42018-09-10 21:33:08 +0300277 itertools_groupby, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000278 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000279};
280
281
282/* _grouper object (internal) ************************************************/
283
284typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000285 PyObject_HEAD
286 PyObject *parent;
287 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000288} _grouperobject;
289
Tal Einat3286ce42018-09-10 21:33:08 +0300290/*[clinic input]
291@classmethod
292itertools._grouper.__new__
293
294 parent: object(subclass_of='&groupby_type')
295 tgtkey: object
296 /
297[clinic start generated code]*/
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000298
299static PyObject *
Tal Einat3286ce42018-09-10 21:33:08 +0300300itertools__grouper_impl(PyTypeObject *type, PyObject *parent,
301 PyObject *tgtkey)
302/*[clinic end generated code: output=462efb1cdebb5914 input=dc180d7771fc8c59]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000303{
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000304 return _grouper_create((groupbyobject*) parent, tgtkey);
305}
306
307static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000308_grouper_create(groupbyobject *parent, PyObject *tgtkey)
309{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000310 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000311
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000312 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
313 if (igo == NULL)
314 return NULL;
315 igo->parent = (PyObject *)parent;
316 Py_INCREF(parent);
317 igo->tgtkey = tgtkey;
318 Py_INCREF(tgtkey);
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300319 parent->currgrouper = igo; /* borrowed reference */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000320
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000321 PyObject_GC_Track(igo);
322 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000323}
324
325static void
326_grouper_dealloc(_grouperobject *igo)
327{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000328 PyObject_GC_UnTrack(igo);
329 Py_DECREF(igo->parent);
330 Py_DECREF(igo->tgtkey);
331 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000332}
333
334static int
335_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
336{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000337 Py_VISIT(igo->parent);
338 Py_VISIT(igo->tgtkey);
339 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000340}
341
342static PyObject *
343_grouper_next(_grouperobject *igo)
344{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000345 groupbyobject *gbo = (groupbyobject *)igo->parent;
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300346 PyObject *r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000347 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000348
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300349 if (gbo->currgrouper != igo)
350 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000351 if (gbo->currvalue == NULL) {
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300352 if (groupby_step(gbo) < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000353 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000354 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000355
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000356 assert(gbo->currkey != NULL);
357 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
358 if (rcmp <= 0)
359 /* got any error or current group is end */
360 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000361
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000362 r = gbo->currvalue;
363 gbo->currvalue = NULL;
364 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000365
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000366 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000367}
368
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000369static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530370_grouper_reduce(_grouperobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000371{
Serhiy Storchakabb86bf42018-12-11 08:28:18 +0200372 _Py_IDENTIFIER(iter);
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300373 if (((groupbyobject *)lz->parent)->currgrouper != lz) {
Serhiy Storchakabb86bf42018-12-11 08:28:18 +0200374 return Py_BuildValue("N(())", _PyEval_GetBuiltinId(&PyId_iter));
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300375 }
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700376 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000377}
378
379static PyMethodDef _grouper_methods[] = {
380 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
381 reduce_doc},
382 {NULL, NULL} /* sentinel */
383};
384
385
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000386static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000387 PyVarObject_HEAD_INIT(NULL, 0)
388 "itertools._grouper", /* tp_name */
389 sizeof(_grouperobject), /* tp_basicsize */
390 0, /* tp_itemsize */
391 /* methods */
392 (destructor)_grouper_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200393 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000394 0, /* tp_getattr */
395 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200396 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000397 0, /* tp_repr */
398 0, /* tp_as_number */
399 0, /* tp_as_sequence */
400 0, /* tp_as_mapping */
401 0, /* tp_hash */
402 0, /* tp_call */
403 0, /* tp_str */
404 PyObject_GenericGetAttr, /* tp_getattro */
405 0, /* tp_setattro */
406 0, /* tp_as_buffer */
407 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
408 0, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700409 (traverseproc)_grouper_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000410 0, /* tp_clear */
411 0, /* tp_richcompare */
412 0, /* tp_weaklistoffset */
413 PyObject_SelfIter, /* tp_iter */
414 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000415 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000416 0, /* tp_members */
417 0, /* tp_getset */
418 0, /* tp_base */
419 0, /* tp_dict */
420 0, /* tp_descr_get */
421 0, /* tp_descr_set */
422 0, /* tp_dictoffset */
423 0, /* tp_init */
424 0, /* tp_alloc */
Tal Einat3286ce42018-09-10 21:33:08 +0300425 itertools__grouper, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000426 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000427};
428
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000429
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700430/* tee object and with supporting function and objects ***********************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000431
Raymond Hettingerad983e72003-11-12 14:32:26 +0000432/* The teedataobject pre-allocates space for LINKCELLS number of objects.
433 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000434 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000435 the other structure members including PyHEAD overhead. The larger the
436 value, the less memory overhead per object and the less time spent
437 allocating/deallocating new links. The smaller the number, the less
438 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000439*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000440#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000441
442typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000443 PyObject_HEAD
444 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700445 int numread; /* 0 <= numread <= LINKCELLS */
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300446 int running;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000447 PyObject *nextlink;
448 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000449} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000450
451typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000452 PyObject_HEAD
453 teedataobject *dataobj;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700454 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000455 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000456} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000457
Raymond Hettingerad983e72003-11-12 14:32:26 +0000458static PyTypeObject teedataobject_type;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000459
460static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000461teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000462{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000463 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000464
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000465 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
466 if (tdo == NULL)
467 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000468
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300469 tdo->running = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000470 tdo->numread = 0;
471 tdo->nextlink = NULL;
472 Py_INCREF(it);
473 tdo->it = it;
474 PyObject_GC_Track(tdo);
475 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000476}
477
478static PyObject *
479teedataobject_jumplink(teedataobject *tdo)
480{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000481 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000482 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000483 Py_XINCREF(tdo->nextlink);
484 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000485}
486
487static PyObject *
488teedataobject_getitem(teedataobject *tdo, int i)
489{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000490 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000491
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000492 assert(i < LINKCELLS);
493 if (i < tdo->numread)
494 value = tdo->values[i];
495 else {
496 /* this is the lead iterator, so fetch more data */
497 assert(i == tdo->numread);
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300498 if (tdo->running) {
499 PyErr_SetString(PyExc_RuntimeError,
500 "cannot re-enter the tee iterator");
501 return NULL;
502 }
503 tdo->running = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000504 value = PyIter_Next(tdo->it);
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300505 tdo->running = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000506 if (value == NULL)
507 return NULL;
508 tdo->numread++;
509 tdo->values[i] = value;
510 }
511 Py_INCREF(value);
512 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000513}
514
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000515static int
516teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
517{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000518 int i;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700519
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000520 Py_VISIT(tdo->it);
521 for (i = 0; i < tdo->numread; i++)
522 Py_VISIT(tdo->values[i]);
523 Py_VISIT(tdo->nextlink);
524 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000525}
526
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200527static void
528teedataobject_safe_decref(PyObject *obj)
529{
530 while (obj && Py_TYPE(obj) == &teedataobject_type &&
531 Py_REFCNT(obj) == 1) {
532 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
533 ((teedataobject *)obj)->nextlink = NULL;
534 Py_DECREF(obj);
535 obj = nextlink;
536 }
537 Py_XDECREF(obj);
538}
539
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000540static int
541teedataobject_clear(teedataobject *tdo)
542{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000543 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200544 PyObject *tmp;
545
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000546 Py_CLEAR(tdo->it);
547 for (i=0 ; i<tdo->numread ; i++)
548 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200549 tmp = tdo->nextlink;
550 tdo->nextlink = NULL;
551 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000552 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000553}
554
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000555static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000556teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000557{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000558 PyObject_GC_UnTrack(tdo);
559 teedataobject_clear(tdo);
560 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000561}
562
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000563static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530564teedataobject_reduce(teedataobject *tdo, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000565{
566 int i;
567 /* create a temporary list of already iterated values */
568 PyObject *values = PyList_New(tdo->numread);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700569
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000570 if (!values)
571 return NULL;
572 for (i=0 ; i<tdo->numread ; i++) {
573 Py_INCREF(tdo->values[i]);
574 PyList_SET_ITEM(values, i, tdo->values[i]);
575 }
576 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
577 values,
578 tdo->nextlink ? tdo->nextlink : Py_None);
579}
580
Tal Einatc4bccd32018-09-12 00:49:13 +0300581/*[clinic input]
582@classmethod
583itertools.teedataobject.__new__
584 iterable as it: object
585 values: object(subclass_of='&PyList_Type')
586 next: object
587 /
588Data container common to multiple tee objects.
589[clinic start generated code]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000590
591static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300592itertools_teedataobject_impl(PyTypeObject *type, PyObject *it,
593 PyObject *values, PyObject *next)
594/*[clinic end generated code: output=3343ceb07e08df5e input=be60f2fabd2b72ba]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000595{
596 teedataobject *tdo;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000597 Py_ssize_t i, len;
598
599 assert(type == &teedataobject_type);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000600
601 tdo = (teedataobject *)teedataobject_newinternal(it);
602 if (!tdo)
603 return NULL;
604
605 len = PyList_GET_SIZE(values);
606 if (len > LINKCELLS)
607 goto err;
608 for (i=0; i<len; i++) {
609 tdo->values[i] = PyList_GET_ITEM(values, i);
610 Py_INCREF(tdo->values[i]);
611 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200612 /* len <= LINKCELLS < INT_MAX */
613 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000614
615 if (len == LINKCELLS) {
616 if (next != Py_None) {
617 if (Py_TYPE(next) != &teedataobject_type)
618 goto err;
619 assert(tdo->nextlink == NULL);
620 Py_INCREF(next);
621 tdo->nextlink = next;
622 }
623 } else {
624 if (next != Py_None)
625 goto err; /* shouldn't have a next if we are not full */
626 }
627 return (PyObject*)tdo;
628
629err:
630 Py_XDECREF(tdo);
631 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
632 return NULL;
633}
634
635static PyMethodDef teedataobject_methods[] = {
636 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
637 reduce_doc},
638 {NULL, NULL} /* sentinel */
639};
640
Raymond Hettingerad983e72003-11-12 14:32:26 +0000641static PyTypeObject teedataobject_type = {
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700642 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000643 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000644 sizeof(teedataobject), /* tp_basicsize */
645 0, /* tp_itemsize */
646 /* methods */
647 (destructor)teedataobject_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200648 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000649 0, /* tp_getattr */
650 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200651 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000652 0, /* tp_repr */
653 0, /* tp_as_number */
654 0, /* tp_as_sequence */
655 0, /* tp_as_mapping */
656 0, /* tp_hash */
657 0, /* tp_call */
658 0, /* tp_str */
659 PyObject_GenericGetAttr, /* tp_getattro */
660 0, /* tp_setattro */
661 0, /* tp_as_buffer */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700662 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300663 itertools_teedataobject__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000664 (traverseproc)teedataobject_traverse, /* tp_traverse */
665 (inquiry)teedataobject_clear, /* tp_clear */
666 0, /* tp_richcompare */
667 0, /* tp_weaklistoffset */
668 0, /* tp_iter */
669 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000670 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000671 0, /* tp_members */
672 0, /* tp_getset */
673 0, /* tp_base */
674 0, /* tp_dict */
675 0, /* tp_descr_get */
676 0, /* tp_descr_set */
677 0, /* tp_dictoffset */
678 0, /* tp_init */
679 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +0300680 itertools_teedataobject, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000681 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000682};
683
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000684
685static PyTypeObject tee_type;
686
687static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000688tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000689{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000690 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000691
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000692 if (to->index >= LINKCELLS) {
693 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200694 if (link == NULL)
695 return NULL;
Serhiy Storchaka57a01d32016-04-10 18:05:40 +0300696 Py_SETREF(to->dataobj, (teedataobject *)link);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000697 to->index = 0;
698 }
699 value = teedataobject_getitem(to->dataobj, to->index);
700 if (value == NULL)
701 return NULL;
702 to->index++;
703 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000704}
705
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000706static int
707tee_traverse(teeobject *to, visitproc visit, void *arg)
708{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000709 Py_VISIT((PyObject *)to->dataobj);
710 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000711}
712
Raymond Hettingerad983e72003-11-12 14:32:26 +0000713static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530714tee_copy(teeobject *to, PyObject *Py_UNUSED(ignored))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000715{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000716 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000717
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000718 newto = PyObject_GC_New(teeobject, &tee_type);
719 if (newto == NULL)
720 return NULL;
721 Py_INCREF(to->dataobj);
722 newto->dataobj = to->dataobj;
723 newto->index = to->index;
724 newto->weakreflist = NULL;
725 PyObject_GC_Track(newto);
726 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000727}
728
729PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
730
731static PyObject *
732tee_fromiterable(PyObject *iterable)
733{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000734 teeobject *to;
735 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000736
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000737 it = PyObject_GetIter(iterable);
738 if (it == NULL)
739 return NULL;
740 if (PyObject_TypeCheck(it, &tee_type)) {
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530741 to = (teeobject *)tee_copy((teeobject *)it, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000742 goto done;
743 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000744
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000745 to = PyObject_GC_New(teeobject, &tee_type);
746 if (to == NULL)
747 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000748 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000749 if (!to->dataobj) {
750 PyObject_GC_Del(to);
751 to = NULL;
752 goto done;
753 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000754
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000755 to->index = 0;
756 to->weakreflist = NULL;
757 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000758done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000759 Py_XDECREF(it);
760 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000761}
762
Tal Einatc4bccd32018-09-12 00:49:13 +0300763/*[clinic input]
764@classmethod
765itertools._tee.__new__
766 iterable: object
767 /
768Iterator wrapped to make it copyable.
769[clinic start generated code]*/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000770
Tal Einatc4bccd32018-09-12 00:49:13 +0300771static PyObject *
772itertools__tee_impl(PyTypeObject *type, PyObject *iterable)
773/*[clinic end generated code: output=b02d3fd26c810c3f input=adc0779d2afe37a2]*/
774{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000775 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000776}
777
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000778static int
779tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000780{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000781 if (to->weakreflist != NULL)
782 PyObject_ClearWeakRefs((PyObject *) to);
783 Py_CLEAR(to->dataobj);
784 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000785}
786
787static void
788tee_dealloc(teeobject *to)
789{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000790 PyObject_GC_UnTrack(to);
791 tee_clear(to);
792 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000793}
794
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000795static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530796tee_reduce(teeobject *to, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000797{
798 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
799}
800
801static PyObject *
802tee_setstate(teeobject *to, PyObject *state)
803{
804 teedataobject *tdo;
805 int index;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300806 if (!PyTuple_Check(state)) {
807 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000808 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300809 }
810 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {
811 return NULL;
812 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000813 if (index < 0 || index > LINKCELLS) {
814 PyErr_SetString(PyExc_ValueError, "Index out of range");
815 return NULL;
816 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200817 Py_INCREF(tdo);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300818 Py_XSETREF(to->dataobj, tdo);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000819 to->index = index;
820 Py_RETURN_NONE;
821}
822
Raymond Hettingerad983e72003-11-12 14:32:26 +0000823static PyMethodDef tee_methods[] = {
Raymond Hettingera6a2d442015-08-16 14:38:07 -0700824 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
825 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
826 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000827 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000828};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000829
830static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000831 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000832 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000833 sizeof(teeobject), /* tp_basicsize */
834 0, /* tp_itemsize */
835 /* methods */
836 (destructor)tee_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200837 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000838 0, /* tp_getattr */
839 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200840 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000841 0, /* tp_repr */
842 0, /* tp_as_number */
843 0, /* tp_as_sequence */
844 0, /* tp_as_mapping */
845 0, /* tp_hash */
846 0, /* tp_call */
847 0, /* tp_str */
848 0, /* tp_getattro */
849 0, /* tp_setattro */
850 0, /* tp_as_buffer */
851 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300852 itertools__tee__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000853 (traverseproc)tee_traverse, /* tp_traverse */
854 (inquiry)tee_clear, /* tp_clear */
855 0, /* tp_richcompare */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700856 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000857 PyObject_SelfIter, /* tp_iter */
858 (iternextfunc)tee_next, /* tp_iternext */
859 tee_methods, /* tp_methods */
860 0, /* tp_members */
861 0, /* tp_getset */
862 0, /* tp_base */
863 0, /* tp_dict */
864 0, /* tp_descr_get */
865 0, /* tp_descr_set */
866 0, /* tp_dictoffset */
867 0, /* tp_init */
868 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +0300869 itertools__tee, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000870 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000871};
872
Tal Einatc4bccd32018-09-12 00:49:13 +0300873/*[clinic input]
874itertools.tee
875 iterable: object
876 n: Py_ssize_t = 2
877 /
878Returns a tuple of n independent iterators.
879[clinic start generated code]*/
880
Raymond Hettingerad983e72003-11-12 14:32:26 +0000881static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300882itertools_tee_impl(PyObject *module, PyObject *iterable, Py_ssize_t n)
883/*[clinic end generated code: output=1c64519cd859c2f0 input=c99a1472c425d66d]*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000884{
Tal Einatc4bccd32018-09-12 00:49:13 +0300885 Py_ssize_t i;
886 PyObject *it, *copyable, *copyfunc, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200887 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000888
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000889 if (n < 0) {
890 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
891 return NULL;
892 }
893 result = PyTuple_New(n);
894 if (result == NULL)
895 return NULL;
896 if (n == 0)
897 return result;
898 it = PyObject_GetIter(iterable);
899 if (it == NULL) {
900 Py_DECREF(result);
901 return NULL;
902 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200903
Serhiy Storchakaf320be72018-01-25 10:49:40 +0200904 if (_PyObject_LookupAttrId(it, &PyId___copy__, &copyfunc) < 0) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200905 Py_DECREF(it);
906 Py_DECREF(result);
907 return NULL;
908 }
Serhiy Storchakaf320be72018-01-25 10:49:40 +0200909 if (copyfunc != NULL) {
910 copyable = it;
911 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200912 else {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000913 copyable = tee_fromiterable(it);
914 Py_DECREF(it);
915 if (copyable == NULL) {
916 Py_DECREF(result);
917 return NULL;
918 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200919 copyfunc = _PyObject_GetAttrId(copyable, &PyId___copy__);
920 if (copyfunc == NULL) {
921 Py_DECREF(copyable);
922 Py_DECREF(result);
923 return NULL;
924 }
925 }
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200926
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200927 PyTuple_SET_ITEM(result, 0, copyable);
928 for (i = 1; i < n; i++) {
929 copyable = _PyObject_CallNoArg(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000930 if (copyable == NULL) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200931 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000932 Py_DECREF(result);
933 return NULL;
934 }
935 PyTuple_SET_ITEM(result, i, copyable);
936 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200937 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000938 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000939}
940
Raymond Hettingerad983e72003-11-12 14:32:26 +0000941
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700942/* cycle object **************************************************************/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000943
944typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000945 PyObject_HEAD
946 PyObject *it;
947 PyObject *saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700948 Py_ssize_t index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000949 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000950} cycleobject;
951
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000952static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000953
Tal Einatc4bccd32018-09-12 00:49:13 +0300954/*[clinic input]
955@classmethod
956itertools.cycle.__new__
957 iterable: object
958 /
959Return elements from the iterable until it is exhausted. Then repeat the sequence indefinitely.
960[clinic start generated code]*/
961
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000962static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300963itertools_cycle_impl(PyTypeObject *type, PyObject *iterable)
964/*[clinic end generated code: output=f60e5ec17a45b35c input=9d1d84bcf66e908b]*/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000965{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000966 PyObject *it;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000967 PyObject *saved;
968 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000969
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000970 /* Get iterator. */
971 it = PyObject_GetIter(iterable);
972 if (it == NULL)
973 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000974
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000975 saved = PyList_New(0);
976 if (saved == NULL) {
977 Py_DECREF(it);
978 return NULL;
979 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000980
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000981 /* create cycleobject structure */
982 lz = (cycleobject *)type->tp_alloc(type, 0);
983 if (lz == NULL) {
984 Py_DECREF(it);
985 Py_DECREF(saved);
986 return NULL;
987 }
988 lz->it = it;
989 lz->saved = saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700990 lz->index = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000991 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000992
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000993 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000994}
995
996static void
997cycle_dealloc(cycleobject *lz)
998{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000999 PyObject_GC_UnTrack(lz);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001000 Py_XDECREF(lz->it);
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001001 Py_XDECREF(lz->saved);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001002 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001003}
1004
1005static int
1006cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
1007{
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001008 if (lz->it)
1009 Py_VISIT(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001010 Py_VISIT(lz->saved);
1011 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001012}
1013
1014static PyObject *
1015cycle_next(cycleobject *lz)
1016{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001017 PyObject *item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001018
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001019 if (lz->it != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001020 item = PyIter_Next(lz->it);
1021 if (item != NULL) {
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001022 if (lz->firstpass)
1023 return item;
1024 if (PyList_Append(lz->saved, item)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001025 Py_DECREF(item);
1026 return NULL;
1027 }
1028 return item;
1029 }
Raymond Hettinger98958fe2015-08-15 15:09:30 -07001030 /* Note: StopIteration is already cleared by PyIter_Next() */
1031 if (PyErr_Occurred())
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001032 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001033 Py_CLEAR(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001034 }
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001035 if (PyList_GET_SIZE(lz->saved) == 0)
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001036 return NULL;
1037 item = PyList_GET_ITEM(lz->saved, lz->index);
1038 lz->index++;
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001039 if (lz->index >= PyList_GET_SIZE(lz->saved))
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001040 lz->index = 0;
1041 Py_INCREF(item);
1042 return item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001043}
1044
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001045static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301046cycle_reduce(cycleobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001047{
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001048 /* Create a new cycle with the iterator tuple, then set the saved state */
1049 if (lz->it == NULL) {
1050 PyObject *it = PyObject_GetIter(lz->saved);
1051 if (it == NULL)
1052 return NULL;
1053 if (lz->index != 0) {
1054 _Py_IDENTIFIER(__setstate__);
1055 PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
1056 "n", lz->index);
1057 if (res == NULL) {
1058 Py_DECREF(it);
1059 return NULL;
1060 }
1061 Py_DECREF(res);
1062 }
Serhiy Storchaka1f21eaa2019-09-01 12:16:51 +03001063 return Py_BuildValue("O(N)(OO)", Py_TYPE(lz), it, lz->saved, Py_True);
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001064 }
Serhiy Storchaka1f21eaa2019-09-01 12:16:51 +03001065 return Py_BuildValue("O(O)(OO)", Py_TYPE(lz), lz->it, lz->saved,
1066 lz->firstpass ? Py_True : Py_False);
Raymond Hettingerc39786d2015-08-15 15:16:12 -07001067}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001068
1069static PyObject *
1070cycle_setstate(cycleobject *lz, PyObject *state)
1071{
1072 PyObject *saved=NULL;
1073 int firstpass;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001074 if (!PyTuple_Check(state)) {
1075 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001076 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001077 }
1078 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
1079 return NULL;
1080 }
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001081 Py_INCREF(saved);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001082 Py_XSETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001083 lz->firstpass = firstpass != 0;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001084 lz->index = 0;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001085 Py_RETURN_NONE;
1086}
1087
1088static PyMethodDef cycle_methods[] = {
1089 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
1090 reduce_doc},
1091 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
1092 setstate_doc},
1093 {NULL, NULL} /* sentinel */
1094};
1095
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001096static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001097 PyVarObject_HEAD_INIT(NULL, 0)
1098 "itertools.cycle", /* tp_name */
1099 sizeof(cycleobject), /* tp_basicsize */
1100 0, /* tp_itemsize */
1101 /* methods */
1102 (destructor)cycle_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001103 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001104 0, /* tp_getattr */
1105 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001106 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001107 0, /* tp_repr */
1108 0, /* tp_as_number */
1109 0, /* tp_as_sequence */
1110 0, /* tp_as_mapping */
1111 0, /* tp_hash */
1112 0, /* tp_call */
1113 0, /* tp_str */
1114 PyObject_GenericGetAttr, /* tp_getattro */
1115 0, /* tp_setattro */
1116 0, /* tp_as_buffer */
1117 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1118 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001119 itertools_cycle__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001120 (traverseproc)cycle_traverse, /* tp_traverse */
1121 0, /* tp_clear */
1122 0, /* tp_richcompare */
1123 0, /* tp_weaklistoffset */
1124 PyObject_SelfIter, /* tp_iter */
1125 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001126 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001127 0, /* tp_members */
1128 0, /* tp_getset */
1129 0, /* tp_base */
1130 0, /* tp_dict */
1131 0, /* tp_descr_get */
1132 0, /* tp_descr_set */
1133 0, /* tp_dictoffset */
1134 0, /* tp_init */
1135 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001136 itertools_cycle, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001137 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001138};
1139
1140
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001141/* dropwhile object **********************************************************/
1142
1143typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001144 PyObject_HEAD
1145 PyObject *func;
1146 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001147 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001148} dropwhileobject;
1149
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001150static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001151
Tal Einatc4bccd32018-09-12 00:49:13 +03001152/*[clinic input]
1153@classmethod
1154itertools.dropwhile.__new__
1155 predicate as func: object
1156 iterable as seq: object
1157 /
1158Drop items from the iterable while predicate(item) is true.
1159
1160Afterwards, return every element until the iterable is exhausted.
1161[clinic start generated code]*/
1162
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001163static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001164itertools_dropwhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1165/*[clinic end generated code: output=92f9d0d89af149e4 input=d39737147c9f0a26]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001166{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001167 PyObject *it;
1168 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001169
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001170 /* Get iterator. */
1171 it = PyObject_GetIter(seq);
1172 if (it == NULL)
1173 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001174
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001175 /* create dropwhileobject structure */
1176 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1177 if (lz == NULL) {
1178 Py_DECREF(it);
1179 return NULL;
1180 }
1181 Py_INCREF(func);
1182 lz->func = func;
1183 lz->it = it;
1184 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001185
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001186 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001187}
1188
1189static void
1190dropwhile_dealloc(dropwhileobject *lz)
1191{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001192 PyObject_GC_UnTrack(lz);
1193 Py_XDECREF(lz->func);
1194 Py_XDECREF(lz->it);
1195 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196}
1197
1198static int
1199dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1200{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001201 Py_VISIT(lz->it);
1202 Py_VISIT(lz->func);
1203 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001204}
1205
1206static PyObject *
1207dropwhile_next(dropwhileobject *lz)
1208{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001209 PyObject *item, *good;
1210 PyObject *it = lz->it;
1211 long ok;
1212 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001213
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001214 iternext = *Py_TYPE(it)->tp_iternext;
1215 for (;;) {
1216 item = iternext(it);
1217 if (item == NULL)
1218 return NULL;
1219 if (lz->start == 1)
1220 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001221
Jeroen Demeyer196a5302019-07-04 12:31:34 +02001222 good = _PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001223 if (good == NULL) {
1224 Py_DECREF(item);
1225 return NULL;
1226 }
1227 ok = PyObject_IsTrue(good);
1228 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001229 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001230 lz->start = 1;
1231 return item;
1232 }
1233 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001234 if (ok < 0)
1235 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001236 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001237}
1238
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001239static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301240dropwhile_reduce(dropwhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001241{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001242 return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->start);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001243}
1244
1245static PyObject *
1246dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1247{
1248 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001249 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001250 return NULL;
1251 lz->start = start;
1252 Py_RETURN_NONE;
1253}
1254
1255static PyMethodDef dropwhile_methods[] = {
1256 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1257 reduce_doc},
1258 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1259 setstate_doc},
1260 {NULL, NULL} /* sentinel */
1261};
1262
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001263static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001264 PyVarObject_HEAD_INIT(NULL, 0)
1265 "itertools.dropwhile", /* tp_name */
1266 sizeof(dropwhileobject), /* tp_basicsize */
1267 0, /* tp_itemsize */
1268 /* methods */
1269 (destructor)dropwhile_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001270 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001271 0, /* tp_getattr */
1272 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001273 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001274 0, /* tp_repr */
1275 0, /* tp_as_number */
1276 0, /* tp_as_sequence */
1277 0, /* tp_as_mapping */
1278 0, /* tp_hash */
1279 0, /* tp_call */
1280 0, /* tp_str */
1281 PyObject_GenericGetAttr, /* tp_getattro */
1282 0, /* tp_setattro */
1283 0, /* tp_as_buffer */
1284 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1285 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001286 itertools_dropwhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001287 (traverseproc)dropwhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001288 0, /* tp_clear */
1289 0, /* tp_richcompare */
1290 0, /* tp_weaklistoffset */
1291 PyObject_SelfIter, /* tp_iter */
1292 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001293 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001294 0, /* tp_members */
1295 0, /* tp_getset */
1296 0, /* tp_base */
1297 0, /* tp_dict */
1298 0, /* tp_descr_get */
1299 0, /* tp_descr_set */
1300 0, /* tp_dictoffset */
1301 0, /* tp_init */
1302 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001303 itertools_dropwhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001304 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001305};
1306
1307
1308/* takewhile object **********************************************************/
1309
1310typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001311 PyObject_HEAD
1312 PyObject *func;
1313 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001314 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001315} takewhileobject;
1316
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001317static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001318
Tal Einatc4bccd32018-09-12 00:49:13 +03001319/*[clinic input]
1320@classmethod
1321itertools.takewhile.__new__
1322 predicate as func: object
1323 iterable as seq: object
1324 /
1325Return successive entries from an iterable as long as the predicate evaluates to true for each entry.
1326[clinic start generated code]*/
1327
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001328static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001329itertools_takewhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1330/*[clinic end generated code: output=bb179ea7864e2ef6 input=ba5255f7519aa119]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001331{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001332 PyObject *it;
1333 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001334
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001335 /* Get iterator. */
1336 it = PyObject_GetIter(seq);
1337 if (it == NULL)
1338 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001339
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001340 /* create takewhileobject structure */
1341 lz = (takewhileobject *)type->tp_alloc(type, 0);
1342 if (lz == NULL) {
1343 Py_DECREF(it);
1344 return NULL;
1345 }
1346 Py_INCREF(func);
1347 lz->func = func;
1348 lz->it = it;
1349 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001350
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001351 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001352}
1353
1354static void
1355takewhile_dealloc(takewhileobject *lz)
1356{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001357 PyObject_GC_UnTrack(lz);
1358 Py_XDECREF(lz->func);
1359 Py_XDECREF(lz->it);
1360 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001361}
1362
1363static int
1364takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1365{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001366 Py_VISIT(lz->it);
1367 Py_VISIT(lz->func);
1368 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001369}
1370
1371static PyObject *
1372takewhile_next(takewhileobject *lz)
1373{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001374 PyObject *item, *good;
1375 PyObject *it = lz->it;
1376 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001377
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001378 if (lz->stop == 1)
1379 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001380
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001381 item = (*Py_TYPE(it)->tp_iternext)(it);
1382 if (item == NULL)
1383 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001384
Jeroen Demeyer196a5302019-07-04 12:31:34 +02001385 good = _PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001386 if (good == NULL) {
1387 Py_DECREF(item);
1388 return NULL;
1389 }
1390 ok = PyObject_IsTrue(good);
1391 Py_DECREF(good);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001392 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001393 return item;
1394 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001395 if (ok == 0)
1396 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001397 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001398}
1399
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001400static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301401takewhile_reduce(takewhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001402{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001403 return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->stop);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001404}
1405
1406static PyObject *
1407takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1408{
1409 int stop = PyObject_IsTrue(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001410
Antoine Pitrou721738f2012-08-15 23:20:39 +02001411 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001412 return NULL;
1413 lz->stop = stop;
1414 Py_RETURN_NONE;
1415}
1416
1417static PyMethodDef takewhile_reduce_methods[] = {
1418 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1419 reduce_doc},
1420 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1421 setstate_doc},
1422 {NULL, NULL} /* sentinel */
1423};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001424
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001425static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001426 PyVarObject_HEAD_INIT(NULL, 0)
1427 "itertools.takewhile", /* tp_name */
1428 sizeof(takewhileobject), /* tp_basicsize */
1429 0, /* tp_itemsize */
1430 /* methods */
1431 (destructor)takewhile_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001432 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001433 0, /* tp_getattr */
1434 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001435 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001436 0, /* tp_repr */
1437 0, /* tp_as_number */
1438 0, /* tp_as_sequence */
1439 0, /* tp_as_mapping */
1440 0, /* tp_hash */
1441 0, /* tp_call */
1442 0, /* tp_str */
1443 PyObject_GenericGetAttr, /* tp_getattro */
1444 0, /* tp_setattro */
1445 0, /* tp_as_buffer */
1446 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1447 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001448 itertools_takewhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001449 (traverseproc)takewhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001450 0, /* tp_clear */
1451 0, /* tp_richcompare */
1452 0, /* tp_weaklistoffset */
1453 PyObject_SelfIter, /* tp_iter */
1454 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001455 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001456 0, /* tp_members */
1457 0, /* tp_getset */
1458 0, /* tp_base */
1459 0, /* tp_dict */
1460 0, /* tp_descr_get */
1461 0, /* tp_descr_set */
1462 0, /* tp_dictoffset */
1463 0, /* tp_init */
1464 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001465 itertools_takewhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001466 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001467};
1468
1469
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001470/* islice object *************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001471
1472typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001473 PyObject_HEAD
1474 PyObject *it;
1475 Py_ssize_t next;
1476 Py_ssize_t stop;
1477 Py_ssize_t step;
1478 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001479} isliceobject;
1480
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001481static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001482
1483static PyObject *
1484islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1485{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001486 PyObject *seq;
1487 Py_ssize_t start=0, stop=-1, step=1;
1488 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1489 Py_ssize_t numargs;
1490 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001491
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001492 if (type == &islice_type && !_PyArg_NoKeywords("islice", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001493 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001494
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001495 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1496 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001497
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001498 numargs = PyTuple_Size(args);
1499 if (numargs == 2) {
1500 if (a1 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001501 stop = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001502 if (stop == -1) {
1503 if (PyErr_Occurred())
1504 PyErr_Clear();
1505 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001506 "Stop argument for islice() must be None or "
1507 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001508 return NULL;
1509 }
1510 }
1511 } else {
1512 if (a1 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001513 start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001514 if (start == -1 && PyErr_Occurred())
1515 PyErr_Clear();
1516 if (a2 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001517 stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001518 if (stop == -1) {
1519 if (PyErr_Occurred())
1520 PyErr_Clear();
1521 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001522 "Stop argument for islice() must be None or "
1523 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001524 return NULL;
1525 }
1526 }
1527 }
1528 if (start<0 || stop<-1) {
1529 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001530 "Indices for islice() must be None or "
1531 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001532 return NULL;
1533 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001534
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001535 if (a3 != NULL) {
1536 if (a3 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001537 step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001538 if (step == -1 && PyErr_Occurred())
1539 PyErr_Clear();
1540 }
1541 if (step<1) {
1542 PyErr_SetString(PyExc_ValueError,
1543 "Step for islice() must be a positive integer or None.");
1544 return NULL;
1545 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001546
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001547 /* Get iterator. */
1548 it = PyObject_GetIter(seq);
1549 if (it == NULL)
1550 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001551
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001552 /* create isliceobject structure */
1553 lz = (isliceobject *)type->tp_alloc(type, 0);
1554 if (lz == NULL) {
1555 Py_DECREF(it);
1556 return NULL;
1557 }
1558 lz->it = it;
1559 lz->next = start;
1560 lz->stop = stop;
1561 lz->step = step;
1562 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001563
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001564 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001565}
1566
1567static void
1568islice_dealloc(isliceobject *lz)
1569{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001570 PyObject_GC_UnTrack(lz);
1571 Py_XDECREF(lz->it);
1572 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001573}
1574
1575static int
1576islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1577{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001578 Py_VISIT(lz->it);
1579 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001580}
1581
1582static PyObject *
1583islice_next(isliceobject *lz)
1584{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001585 PyObject *item;
1586 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001587 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001588 Py_ssize_t oldnext;
1589 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001590
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001591 if (it == NULL)
1592 return NULL;
1593
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001594 iternext = *Py_TYPE(it)->tp_iternext;
1595 while (lz->cnt < lz->next) {
1596 item = iternext(it);
1597 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001598 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001599 Py_DECREF(item);
1600 lz->cnt++;
1601 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001602 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001603 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001604 item = iternext(it);
1605 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001606 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001607 lz->cnt++;
1608 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001609 /* The (size_t) cast below avoids the danger of undefined
1610 behaviour from signed integer overflow. */
1611 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001612 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1613 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001614 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001615
1616empty:
1617 Py_CLEAR(lz->it);
1618 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001619}
1620
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001621static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301622islice_reduce(isliceobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001623{
1624 /* When unpickled, generate a new object with the same bounds,
1625 * then 'setstate' with the next and count
1626 */
1627 PyObject *stop;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001628
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001629 if (lz->it == NULL) {
1630 PyObject *empty_list;
1631 PyObject *empty_it;
1632 empty_list = PyList_New(0);
1633 if (empty_list == NULL)
1634 return NULL;
1635 empty_it = PyObject_GetIter(empty_list);
1636 Py_DECREF(empty_list);
1637 if (empty_it == NULL)
1638 return NULL;
1639 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1640 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001641 if (lz->stop == -1) {
1642 stop = Py_None;
1643 Py_INCREF(stop);
1644 } else {
1645 stop = PyLong_FromSsize_t(lz->stop);
1646 if (stop == NULL)
1647 return NULL;
1648 }
1649 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1650 lz->it, lz->next, stop, lz->step,
1651 lz->cnt);
1652}
1653
1654static PyObject *
1655islice_setstate(isliceobject *lz, PyObject *state)
1656{
1657 Py_ssize_t cnt = PyLong_AsSsize_t(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001658
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001659 if (cnt == -1 && PyErr_Occurred())
1660 return NULL;
1661 lz->cnt = cnt;
1662 Py_RETURN_NONE;
1663}
1664
1665static PyMethodDef islice_methods[] = {
1666 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1667 reduce_doc},
1668 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1669 setstate_doc},
1670 {NULL, NULL} /* sentinel */
1671};
1672
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001673PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001674"islice(iterable, stop) --> islice object\n\
1675islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001676\n\
1677Return an iterator whose next() method returns selected values from an\n\
1678iterable. If start is specified, will skip all preceding elements;\n\
1679otherwise, start defaults to zero. Step defaults to one. If\n\
oldkaa0735f2018-02-02 16:52:55 +08001680specified as another value, step determines how many values are\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001681skipped between successive calls. Works like a slice() on a list\n\
1682but returns an iterator.");
1683
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001684static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001685 PyVarObject_HEAD_INIT(NULL, 0)
1686 "itertools.islice", /* tp_name */
1687 sizeof(isliceobject), /* tp_basicsize */
1688 0, /* tp_itemsize */
1689 /* methods */
1690 (destructor)islice_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001691 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001692 0, /* tp_getattr */
1693 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001694 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001695 0, /* tp_repr */
1696 0, /* tp_as_number */
1697 0, /* tp_as_sequence */
1698 0, /* tp_as_mapping */
1699 0, /* tp_hash */
1700 0, /* tp_call */
1701 0, /* tp_str */
1702 PyObject_GenericGetAttr, /* tp_getattro */
1703 0, /* tp_setattro */
1704 0, /* tp_as_buffer */
1705 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1706 Py_TPFLAGS_BASETYPE, /* tp_flags */
1707 islice_doc, /* tp_doc */
1708 (traverseproc)islice_traverse, /* tp_traverse */
1709 0, /* tp_clear */
1710 0, /* tp_richcompare */
1711 0, /* tp_weaklistoffset */
1712 PyObject_SelfIter, /* tp_iter */
1713 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001714 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001715 0, /* tp_members */
1716 0, /* tp_getset */
1717 0, /* tp_base */
1718 0, /* tp_dict */
1719 0, /* tp_descr_get */
1720 0, /* tp_descr_set */
1721 0, /* tp_dictoffset */
1722 0, /* tp_init */
1723 0, /* tp_alloc */
1724 islice_new, /* tp_new */
1725 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001726};
1727
1728
1729/* starmap object ************************************************************/
1730
1731typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001732 PyObject_HEAD
1733 PyObject *func;
1734 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001735} starmapobject;
1736
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001737static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001738
Tal Einatc4bccd32018-09-12 00:49:13 +03001739/*[clinic input]
1740@classmethod
1741itertools.starmap.__new__
1742 function as func: object
1743 iterable as seq: object
1744 /
1745Return an iterator whose values are returned from the function evaluated with an argument tuple taken from the given sequence.
1746[clinic start generated code]*/
1747
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001748static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001749itertools_starmap_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1750/*[clinic end generated code: output=79eeb81d452c6e8d input=844766df6a0d4dad]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001751{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001752 PyObject *it;
1753 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001754
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001755 /* Get iterator. */
1756 it = PyObject_GetIter(seq);
1757 if (it == NULL)
1758 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001759
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001760 /* create starmapobject structure */
1761 lz = (starmapobject *)type->tp_alloc(type, 0);
1762 if (lz == NULL) {
1763 Py_DECREF(it);
1764 return NULL;
1765 }
1766 Py_INCREF(func);
1767 lz->func = func;
1768 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001769
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001770 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001771}
1772
1773static void
1774starmap_dealloc(starmapobject *lz)
1775{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001776 PyObject_GC_UnTrack(lz);
1777 Py_XDECREF(lz->func);
1778 Py_XDECREF(lz->it);
1779 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001780}
1781
1782static int
1783starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1784{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001785 Py_VISIT(lz->it);
1786 Py_VISIT(lz->func);
1787 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001788}
1789
1790static PyObject *
1791starmap_next(starmapobject *lz)
1792{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001793 PyObject *args;
1794 PyObject *result;
1795 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001796
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001797 args = (*Py_TYPE(it)->tp_iternext)(it);
1798 if (args == NULL)
1799 return NULL;
1800 if (!PyTuple_CheckExact(args)) {
1801 PyObject *newargs = PySequence_Tuple(args);
1802 Py_DECREF(args);
1803 if (newargs == NULL)
1804 return NULL;
1805 args = newargs;
1806 }
1807 result = PyObject_Call(lz->func, args, NULL);
1808 Py_DECREF(args);
1809 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001810}
1811
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001812static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301813starmap_reduce(starmapobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001814{
1815 /* Just pickle the iterator */
1816 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1817}
1818
1819static PyMethodDef starmap_methods[] = {
1820 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1821 reduce_doc},
1822 {NULL, NULL} /* sentinel */
1823};
1824
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001825static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001826 PyVarObject_HEAD_INIT(NULL, 0)
1827 "itertools.starmap", /* tp_name */
1828 sizeof(starmapobject), /* tp_basicsize */
1829 0, /* tp_itemsize */
1830 /* methods */
1831 (destructor)starmap_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001832 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001833 0, /* tp_getattr */
1834 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001835 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001836 0, /* tp_repr */
1837 0, /* tp_as_number */
1838 0, /* tp_as_sequence */
1839 0, /* tp_as_mapping */
1840 0, /* tp_hash */
1841 0, /* tp_call */
1842 0, /* tp_str */
1843 PyObject_GenericGetAttr, /* tp_getattro */
1844 0, /* tp_setattro */
1845 0, /* tp_as_buffer */
1846 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1847 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001848 itertools_starmap__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001849 (traverseproc)starmap_traverse, /* tp_traverse */
1850 0, /* tp_clear */
1851 0, /* tp_richcompare */
1852 0, /* tp_weaklistoffset */
1853 PyObject_SelfIter, /* tp_iter */
1854 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001855 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001856 0, /* tp_members */
1857 0, /* tp_getset */
1858 0, /* tp_base */
1859 0, /* tp_dict */
1860 0, /* tp_descr_get */
1861 0, /* tp_descr_set */
1862 0, /* tp_dictoffset */
1863 0, /* tp_init */
1864 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001865 itertools_starmap, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001866 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001867};
1868
1869
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001870/* chain object **************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001871
1872typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001873 PyObject_HEAD
1874 PyObject *source; /* Iterator over input iterables */
1875 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001876} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001877
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001878static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001879
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001880static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001881chain_new_internal(PyTypeObject *type, PyObject *source)
1882{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001883 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001884
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001885 lz = (chainobject *)type->tp_alloc(type, 0);
1886 if (lz == NULL) {
1887 Py_DECREF(source);
1888 return NULL;
1889 }
1890
1891 lz->source = source;
1892 lz->active = NULL;
1893 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001894}
1895
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001896static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001897chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001898{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001899 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001900
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001901 if (type == &chain_type && !_PyArg_NoKeywords("chain", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001902 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001903
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001904 source = PyObject_GetIter(args);
1905 if (source == NULL)
1906 return NULL;
1907
1908 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001909}
1910
Tal Einatc4bccd32018-09-12 00:49:13 +03001911/*[clinic input]
1912@classmethod
1913itertools.chain.from_iterable
1914 iterable as arg: object
1915 /
1916Alternative chain() constructor taking a single iterable argument that evaluates lazily.
1917[clinic start generated code]*/
1918
Christian Heimesf16baeb2008-02-29 14:57:44 +00001919static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001920itertools_chain_from_iterable(PyTypeObject *type, PyObject *arg)
1921/*[clinic end generated code: output=667ae7a7f7b68654 input=72c39e3a2ca3be85]*/
Christian Heimesf16baeb2008-02-29 14:57:44 +00001922{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001923 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001924
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001925 source = PyObject_GetIter(arg);
1926 if (source == NULL)
1927 return NULL;
1928
1929 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001930}
1931
1932static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001933chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001934{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001935 PyObject_GC_UnTrack(lz);
1936 Py_XDECREF(lz->active);
1937 Py_XDECREF(lz->source);
1938 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001939}
1940
Raymond Hettinger2012f172003-02-07 05:32:58 +00001941static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001942chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001943{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001944 Py_VISIT(lz->source);
1945 Py_VISIT(lz->active);
1946 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001947}
1948
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001949static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001950chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001951{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001952 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001953
T. Wouters5466d4a2017-03-30 09:58:35 -07001954 /* lz->source is the iterator of iterables. If it's NULL, we've already
1955 * consumed them all. lz->active is the current iterator. If it's NULL,
1956 * we should grab a new one from lz->source. */
1957 while (lz->source != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001958 if (lz->active == NULL) {
T. Wouters5466d4a2017-03-30 09:58:35 -07001959 PyObject *iterable = PyIter_Next(lz->source);
1960 if (iterable == NULL) {
1961 Py_CLEAR(lz->source);
1962 return NULL; /* no more input sources */
1963 }
1964 lz->active = PyObject_GetIter(iterable);
1965 Py_DECREF(iterable);
1966 if (lz->active == NULL) {
1967 Py_CLEAR(lz->source);
1968 return NULL; /* input not iterable */
1969 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001970 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001971 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
1972 if (item != NULL)
1973 return item;
1974 if (PyErr_Occurred()) {
1975 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1976 PyErr_Clear();
1977 else
1978 return NULL; /* input raised an exception */
1979 }
1980 /* lz->active is consumed, try with the next iterable. */
1981 Py_CLEAR(lz->active);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001982 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001983 /* Everything had been consumed already. */
1984 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001985}
1986
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001987static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301988chain_reduce(chainobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001989{
1990 if (lz->source) {
1991 /* we can't pickle function objects (itertools.from_iterable) so
1992 * we must use setstate to replace the iterable. One day we
1993 * will fix pickling of functions
1994 */
1995 if (lz->active) {
1996 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1997 } else {
1998 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1999 }
2000 } else {
2001 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
2002 }
2003 return NULL;
2004}
2005
2006static PyObject *
2007chain_setstate(chainobject *lz, PyObject *state)
2008{
2009 PyObject *source, *active=NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002010
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002011 if (!PyTuple_Check(state)) {
2012 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002013 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002014 }
2015 if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
2016 return NULL;
2017 }
2018 if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
2019 PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
2020 return NULL;
2021 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002022
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002023 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002024 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002025 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002026 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002027 Py_RETURN_NONE;
2028}
2029
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002030PyDoc_STRVAR(chain_doc,
2031"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002032\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002033Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002034first iterable until it is exhausted, then elements from the next\n\
2035iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002036
Christian Heimesf16baeb2008-02-29 14:57:44 +00002037static PyMethodDef chain_methods[] = {
Tal Einatc4bccd32018-09-12 00:49:13 +03002038 ITERTOOLS_CHAIN_FROM_ITERABLE_METHODDEF
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002039 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
2040 reduce_doc},
2041 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
2042 setstate_doc},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07002043 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00002044};
2045
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002046static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002047 PyVarObject_HEAD_INIT(NULL, 0)
2048 "itertools.chain", /* tp_name */
2049 sizeof(chainobject), /* tp_basicsize */
2050 0, /* tp_itemsize */
2051 /* methods */
2052 (destructor)chain_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002053 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002054 0, /* tp_getattr */
2055 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002056 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002057 0, /* tp_repr */
2058 0, /* tp_as_number */
2059 0, /* tp_as_sequence */
2060 0, /* tp_as_mapping */
2061 0, /* tp_hash */
2062 0, /* tp_call */
2063 0, /* tp_str */
2064 PyObject_GenericGetAttr, /* tp_getattro */
2065 0, /* tp_setattro */
2066 0, /* tp_as_buffer */
2067 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2068 Py_TPFLAGS_BASETYPE, /* tp_flags */
2069 chain_doc, /* tp_doc */
2070 (traverseproc)chain_traverse, /* tp_traverse */
2071 0, /* tp_clear */
2072 0, /* tp_richcompare */
2073 0, /* tp_weaklistoffset */
2074 PyObject_SelfIter, /* tp_iter */
2075 (iternextfunc)chain_next, /* tp_iternext */
2076 chain_methods, /* tp_methods */
2077 0, /* tp_members */
2078 0, /* tp_getset */
2079 0, /* tp_base */
2080 0, /* tp_dict */
2081 0, /* tp_descr_get */
2082 0, /* tp_descr_set */
2083 0, /* tp_dictoffset */
2084 0, /* tp_init */
2085 0, /* tp_alloc */
2086 chain_new, /* tp_new */
2087 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002088};
2089
2090
Christian Heimesc3f30c42008-02-22 16:37:40 +00002091/* product object ************************************************************/
2092
2093typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002094 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002095 PyObject *pools; /* tuple of pool tuples */
2096 Py_ssize_t *indices; /* one index per pool */
2097 PyObject *result; /* most recently returned result tuple */
2098 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002099} productobject;
2100
2101static PyTypeObject product_type;
2102
2103static PyObject *
2104product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2105{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002106 productobject *lz;
2107 Py_ssize_t nargs, npools, repeat=1;
2108 PyObject *pools = NULL;
2109 Py_ssize_t *indices = NULL;
2110 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002111
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002112 if (kwds != NULL) {
2113 char *kwlist[] = {"repeat", 0};
2114 PyObject *tmpargs = PyTuple_New(0);
2115 if (tmpargs == NULL)
2116 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002117 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2118 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002119 Py_DECREF(tmpargs);
2120 return NULL;
2121 }
2122 Py_DECREF(tmpargs);
2123 if (repeat < 0) {
2124 PyErr_SetString(PyExc_ValueError,
2125 "repeat argument cannot be negative");
2126 return NULL;
2127 }
2128 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002129
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002130 assert(PyTuple_CheckExact(args));
2131 if (repeat == 0) {
2132 nargs = 0;
2133 } else {
2134 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002135 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002136 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2137 return NULL;
2138 }
2139 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002140 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002141
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002142 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002143 if (indices == NULL) {
2144 PyErr_NoMemory();
2145 goto error;
2146 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002147
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002148 pools = PyTuple_New(npools);
2149 if (pools == NULL)
2150 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002151
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002152 for (i=0; i < nargs ; ++i) {
2153 PyObject *item = PyTuple_GET_ITEM(args, i);
2154 PyObject *pool = PySequence_Tuple(item);
2155 if (pool == NULL)
2156 goto error;
2157 PyTuple_SET_ITEM(pools, i, pool);
2158 indices[i] = 0;
2159 }
2160 for ( ; i < npools; ++i) {
2161 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2162 Py_INCREF(pool);
2163 PyTuple_SET_ITEM(pools, i, pool);
2164 indices[i] = 0;
2165 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002166
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002167 /* create productobject structure */
2168 lz = (productobject *)type->tp_alloc(type, 0);
2169 if (lz == NULL)
2170 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002171
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002172 lz->pools = pools;
2173 lz->indices = indices;
2174 lz->result = NULL;
2175 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002176
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002177 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002178
2179error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002180 if (indices != NULL)
2181 PyMem_Free(indices);
2182 Py_XDECREF(pools);
2183 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002184}
2185
2186static void
2187product_dealloc(productobject *lz)
2188{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002189 PyObject_GC_UnTrack(lz);
2190 Py_XDECREF(lz->pools);
2191 Py_XDECREF(lz->result);
2192 if (lz->indices != NULL)
2193 PyMem_Free(lz->indices);
2194 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002195}
2196
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002197static PyObject *
2198product_sizeof(productobject *lz, void *unused)
2199{
2200 Py_ssize_t res;
2201
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002202 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002203 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2204 return PyLong_FromSsize_t(res);
2205}
2206
2207PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2208
Christian Heimesc3f30c42008-02-22 16:37:40 +00002209static int
2210product_traverse(productobject *lz, visitproc visit, void *arg)
2211{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002212 Py_VISIT(lz->pools);
2213 Py_VISIT(lz->result);
2214 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002215}
2216
2217static PyObject *
2218product_next(productobject *lz)
2219{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002220 PyObject *pool;
2221 PyObject *elem;
2222 PyObject *oldelem;
2223 PyObject *pools = lz->pools;
2224 PyObject *result = lz->result;
2225 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2226 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002227
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002228 if (lz->stopped)
2229 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002230
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002231 if (result == NULL) {
2232 /* On the first pass, return an initial tuple filled with the
2233 first element from each pool. */
2234 result = PyTuple_New(npools);
2235 if (result == NULL)
2236 goto empty;
2237 lz->result = result;
2238 for (i=0; i < npools; i++) {
2239 pool = PyTuple_GET_ITEM(pools, i);
2240 if (PyTuple_GET_SIZE(pool) == 0)
2241 goto empty;
2242 elem = PyTuple_GET_ITEM(pool, 0);
2243 Py_INCREF(elem);
2244 PyTuple_SET_ITEM(result, i, elem);
2245 }
2246 } else {
2247 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002248
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002249 /* Copy the previous result tuple or re-use it if available */
2250 if (Py_REFCNT(result) > 1) {
2251 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002252 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002253 if (result == NULL)
2254 goto empty;
2255 lz->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002256 Py_DECREF(old_result);
2257 }
2258 /* Now, we've got the only copy so we can update it in-place */
2259 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002260
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002261 /* Update the pool indices right-to-left. Only advance to the
2262 next pool when the previous one rolls-over */
2263 for (i=npools-1 ; i >= 0 ; i--) {
2264 pool = PyTuple_GET_ITEM(pools, i);
2265 indices[i]++;
2266 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2267 /* Roll-over and advance to next pool */
2268 indices[i] = 0;
2269 elem = PyTuple_GET_ITEM(pool, 0);
2270 Py_INCREF(elem);
2271 oldelem = PyTuple_GET_ITEM(result, i);
2272 PyTuple_SET_ITEM(result, i, elem);
2273 Py_DECREF(oldelem);
2274 } else {
2275 /* No rollover. Just increment and stop here. */
2276 elem = PyTuple_GET_ITEM(pool, indices[i]);
2277 Py_INCREF(elem);
2278 oldelem = PyTuple_GET_ITEM(result, i);
2279 PyTuple_SET_ITEM(result, i, elem);
2280 Py_DECREF(oldelem);
2281 break;
2282 }
2283 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002284
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002285 /* If i is negative, then the indices have all rolled-over
2286 and we're done. */
2287 if (i < 0)
2288 goto empty;
2289 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002290
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002291 Py_INCREF(result);
2292 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002293
2294empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002295 lz->stopped = 1;
2296 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002297}
2298
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002299static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302300product_reduce(productobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002301{
2302 if (lz->stopped) {
2303 return Py_BuildValue("O(())", Py_TYPE(lz));
2304 } else if (lz->result == NULL) {
2305 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2306 } else {
2307 PyObject *indices;
2308 Py_ssize_t n, i;
2309
2310 /* we must pickle the indices use them for setstate, and
2311 * additionally indicate that the iterator has started
2312 */
2313 n = PyTuple_GET_SIZE(lz->pools);
2314 indices = PyTuple_New(n);
2315 if (indices == NULL)
2316 return NULL;
2317 for (i=0; i<n; i++){
2318 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2319 if (!index) {
2320 Py_DECREF(indices);
2321 return NULL;
2322 }
2323 PyTuple_SET_ITEM(indices, i, index);
2324 }
2325 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2326 }
2327}
2328
2329static PyObject *
2330product_setstate(productobject *lz, PyObject *state)
2331{
2332 PyObject *result;
2333 Py_ssize_t n, i;
2334
2335 n = PyTuple_GET_SIZE(lz->pools);
2336 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2337 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2338 return NULL;
2339 }
2340 for (i=0; i<n; i++)
2341 {
2342 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2343 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002344 PyObject* pool;
2345 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002346 if (index < 0 && PyErr_Occurred())
2347 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002348 pool = PyTuple_GET_ITEM(lz->pools, i);
2349 poolsize = PyTuple_GET_SIZE(pool);
2350 if (poolsize == 0) {
2351 lz->stopped = 1;
2352 Py_RETURN_NONE;
2353 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002354 /* clamp the index */
2355 if (index < 0)
2356 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002357 else if (index > poolsize-1)
2358 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002359 lz->indices[i] = index;
2360 }
2361
2362 result = PyTuple_New(n);
2363 if (!result)
2364 return NULL;
2365 for (i=0; i<n; i++) {
2366 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2367 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2368 Py_INCREF(element);
2369 PyTuple_SET_ITEM(result, i, element);
2370 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002371 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002372 Py_RETURN_NONE;
2373}
2374
2375static PyMethodDef product_methods[] = {
2376 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2377 reduce_doc},
2378 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2379 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002380 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2381 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002382 {NULL, NULL} /* sentinel */
2383};
2384
Christian Heimesc3f30c42008-02-22 16:37:40 +00002385PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002386"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002387\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002388Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002389For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2390The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2391cycle in a manner similar to an odometer (with the rightmost element changing\n\
2392on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002393To compute the product of an iterable with itself, specify the number\n\
2394of repetitions with the optional repeat keyword argument. For example,\n\
2395product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002396product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2397product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2398
2399static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002400 PyVarObject_HEAD_INIT(NULL, 0)
2401 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002402 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002403 0, /* tp_itemsize */
2404 /* methods */
2405 (destructor)product_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002406 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002407 0, /* tp_getattr */
2408 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002409 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002410 0, /* tp_repr */
2411 0, /* tp_as_number */
2412 0, /* tp_as_sequence */
2413 0, /* tp_as_mapping */
2414 0, /* tp_hash */
2415 0, /* tp_call */
2416 0, /* tp_str */
2417 PyObject_GenericGetAttr, /* tp_getattro */
2418 0, /* tp_setattro */
2419 0, /* tp_as_buffer */
2420 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2421 Py_TPFLAGS_BASETYPE, /* tp_flags */
2422 product_doc, /* tp_doc */
2423 (traverseproc)product_traverse, /* tp_traverse */
2424 0, /* tp_clear */
2425 0, /* tp_richcompare */
2426 0, /* tp_weaklistoffset */
2427 PyObject_SelfIter, /* tp_iter */
2428 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002429 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002430 0, /* tp_members */
2431 0, /* tp_getset */
2432 0, /* tp_base */
2433 0, /* tp_dict */
2434 0, /* tp_descr_get */
2435 0, /* tp_descr_set */
2436 0, /* tp_dictoffset */
2437 0, /* tp_init */
2438 0, /* tp_alloc */
2439 product_new, /* tp_new */
2440 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002441};
2442
2443
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002444/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002445
2446typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002447 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002448 PyObject *pool; /* input converted to a tuple */
2449 Py_ssize_t *indices; /* one index per result element */
2450 PyObject *result; /* most recently returned result tuple */
2451 Py_ssize_t r; /* size of result tuple */
2452 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002453} combinationsobject;
2454
2455static PyTypeObject combinations_type;
2456
Tal Einatc4bccd32018-09-12 00:49:13 +03002457
2458/*[clinic input]
2459@classmethod
2460itertools.combinations.__new__
2461 iterable: object
2462 r: Py_ssize_t
2463Return successive r-length combinations of elements in the iterable.
2464
2465combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)
2466[clinic start generated code]*/
2467
Christian Heimes380f7f22008-02-28 11:19:05 +00002468static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002469itertools_combinations_impl(PyTypeObject *type, PyObject *iterable,
2470 Py_ssize_t r)
2471/*[clinic end generated code: output=87a689b39c40039c input=06bede09e3da20f8]*/
Christian Heimes380f7f22008-02-28 11:19:05 +00002472{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002473 combinationsobject *co;
2474 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002475 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002476 Py_ssize_t *indices = NULL;
2477 Py_ssize_t i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002478
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002479 pool = PySequence_Tuple(iterable);
2480 if (pool == NULL)
2481 goto error;
2482 n = PyTuple_GET_SIZE(pool);
2483 if (r < 0) {
2484 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2485 goto error;
2486 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002487
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002488 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002489 if (indices == NULL) {
2490 PyErr_NoMemory();
2491 goto error;
2492 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002493
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002494 for (i=0 ; i<r ; i++)
2495 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002496
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002497 /* create combinationsobject structure */
2498 co = (combinationsobject *)type->tp_alloc(type, 0);
2499 if (co == NULL)
2500 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002501
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002502 co->pool = pool;
2503 co->indices = indices;
2504 co->result = NULL;
2505 co->r = r;
2506 co->stopped = r > n ? 1 : 0;
2507
2508 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002509
2510error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002511 if (indices != NULL)
2512 PyMem_Free(indices);
2513 Py_XDECREF(pool);
2514 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002515}
2516
2517static void
2518combinations_dealloc(combinationsobject *co)
2519{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002520 PyObject_GC_UnTrack(co);
2521 Py_XDECREF(co->pool);
2522 Py_XDECREF(co->result);
2523 if (co->indices != NULL)
2524 PyMem_Free(co->indices);
2525 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002526}
2527
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002528static PyObject *
2529combinations_sizeof(combinationsobject *co, void *unused)
2530{
2531 Py_ssize_t res;
2532
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002533 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002534 res += co->r * sizeof(Py_ssize_t);
2535 return PyLong_FromSsize_t(res);
2536}
2537
Christian Heimes380f7f22008-02-28 11:19:05 +00002538static int
2539combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2540{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002541 Py_VISIT(co->pool);
2542 Py_VISIT(co->result);
2543 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002544}
2545
2546static PyObject *
2547combinations_next(combinationsobject *co)
2548{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002549 PyObject *elem;
2550 PyObject *oldelem;
2551 PyObject *pool = co->pool;
2552 Py_ssize_t *indices = co->indices;
2553 PyObject *result = co->result;
2554 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2555 Py_ssize_t r = co->r;
2556 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002557
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002558 if (co->stopped)
2559 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002560
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002561 if (result == NULL) {
2562 /* On the first pass, initialize result tuple using the indices */
2563 result = PyTuple_New(r);
2564 if (result == NULL)
2565 goto empty;
2566 co->result = result;
2567 for (i=0; i<r ; i++) {
2568 index = indices[i];
2569 elem = PyTuple_GET_ITEM(pool, index);
2570 Py_INCREF(elem);
2571 PyTuple_SET_ITEM(result, i, elem);
2572 }
2573 } else {
2574 /* Copy the previous result tuple or re-use it if available */
2575 if (Py_REFCNT(result) > 1) {
2576 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002577 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002578 if (result == NULL)
2579 goto empty;
2580 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002581 Py_DECREF(old_result);
2582 }
2583 /* Now, we've got the only copy so we can update it in-place
2584 * CPython's empty tuple is a singleton and cached in
2585 * PyTuple's freelist.
2586 */
2587 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002588
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002589 /* Scan indices right-to-left until finding one that is not
2590 at its maximum (i + n - r). */
2591 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2592 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002593
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002594 /* If i is negative, then the indices are all at
2595 their maximum value and we're done. */
2596 if (i < 0)
2597 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002598
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002599 /* Increment the current index which we know is not at its
2600 maximum. Then move back to the right setting each index
2601 to its lowest possible value (one higher than the index
2602 to its left -- this maintains the sort order invariant). */
2603 indices[i]++;
2604 for (j=i+1 ; j<r ; j++)
2605 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002606
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002607 /* Update the result tuple for the new indices
2608 starting with i, the leftmost index that changed */
2609 for ( ; i<r ; i++) {
2610 index = indices[i];
2611 elem = PyTuple_GET_ITEM(pool, index);
2612 Py_INCREF(elem);
2613 oldelem = PyTuple_GET_ITEM(result, i);
2614 PyTuple_SET_ITEM(result, i, elem);
2615 Py_DECREF(oldelem);
2616 }
2617 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002618
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002619 Py_INCREF(result);
2620 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002621
2622empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002623 co->stopped = 1;
2624 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002625}
2626
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002627static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302628combinations_reduce(combinationsobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002629{
2630 if (lz->result == NULL) {
2631 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2632 } else if (lz->stopped) {
2633 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2634 } else {
2635 PyObject *indices;
2636 Py_ssize_t i;
2637
2638 /* we must pickle the indices and use them for setstate */
2639 indices = PyTuple_New(lz->r);
2640 if (!indices)
2641 return NULL;
2642 for (i=0; i<lz->r; i++)
2643 {
2644 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2645 if (!index) {
2646 Py_DECREF(indices);
2647 return NULL;
2648 }
2649 PyTuple_SET_ITEM(indices, i, index);
2650 }
2651
2652 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2653 }
2654}
2655
2656static PyObject *
2657combinations_setstate(combinationsobject *lz, PyObject *state)
2658{
2659 PyObject *result;
2660 Py_ssize_t i;
2661 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2662
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002663 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002664 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2665 return NULL;
2666 }
2667
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002668 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002669 Py_ssize_t max;
2670 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2671 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002672
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002673 if (index == -1 && PyErr_Occurred())
2674 return NULL; /* not an integer */
2675 max = i + n - lz->r;
2676 /* clamp the index (beware of negative max) */
2677 if (index > max)
2678 index = max;
2679 if (index < 0)
2680 index = 0;
2681 lz->indices[i] = index;
2682 }
2683
2684 result = PyTuple_New(lz->r);
2685 if (result == NULL)
2686 return NULL;
2687 for (i=0; i<lz->r; i++) {
2688 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2689 Py_INCREF(element);
2690 PyTuple_SET_ITEM(result, i, element);
2691 }
2692
Serhiy Storchaka48842712016-04-06 09:45:48 +03002693 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002694 Py_RETURN_NONE;
2695}
2696
2697static PyMethodDef combinations_methods[] = {
2698 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2699 reduce_doc},
2700 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2701 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002702 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2703 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002704 {NULL, NULL} /* sentinel */
2705};
2706
Christian Heimes380f7f22008-02-28 11:19:05 +00002707static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002708 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002709 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002710 sizeof(combinationsobject), /* tp_basicsize */
2711 0, /* tp_itemsize */
2712 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002713 (destructor)combinations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002714 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002715 0, /* tp_getattr */
2716 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002717 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002718 0, /* tp_repr */
2719 0, /* tp_as_number */
2720 0, /* tp_as_sequence */
2721 0, /* tp_as_mapping */
2722 0, /* tp_hash */
2723 0, /* tp_call */
2724 0, /* tp_str */
2725 PyObject_GenericGetAttr, /* tp_getattro */
2726 0, /* tp_setattro */
2727 0, /* tp_as_buffer */
2728 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2729 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03002730 itertools_combinations__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002731 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002732 0, /* tp_clear */
2733 0, /* tp_richcompare */
2734 0, /* tp_weaklistoffset */
2735 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002736 (iternextfunc)combinations_next, /* tp_iternext */
2737 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002738 0, /* tp_members */
2739 0, /* tp_getset */
2740 0, /* tp_base */
2741 0, /* tp_dict */
2742 0, /* tp_descr_get */
2743 0, /* tp_descr_set */
2744 0, /* tp_dictoffset */
2745 0, /* tp_init */
2746 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03002747 itertools_combinations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002748 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002749};
2750
2751
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002752/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002753
2754/* Equivalent to:
2755
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002756 def combinations_with_replacement(iterable, r):
2757 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2758 # number items returned: (n+r-1)! / r! / (n-1)!
2759 pool = tuple(iterable)
2760 n = len(pool)
2761 indices = [0] * r
2762 yield tuple(pool[i] for i in indices)
2763 while 1:
2764 for i in reversed(range(r)):
2765 if indices[i] != n - 1:
2766 break
2767 else:
2768 return
2769 indices[i:] = [indices[i] + 1] * (r - i)
2770 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002771
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002772 def combinations_with_replacement2(iterable, r):
2773 'Alternate version that filters from product()'
2774 pool = tuple(iterable)
2775 n = len(pool)
2776 for indices in product(range(n), repeat=r):
2777 if sorted(indices) == list(indices):
2778 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002779*/
2780typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002781 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002782 PyObject *pool; /* input converted to a tuple */
2783 Py_ssize_t *indices; /* one index per result element */
2784 PyObject *result; /* most recently returned result tuple */
2785 Py_ssize_t r; /* size of result tuple */
2786 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002787} cwrobject;
2788
2789static PyTypeObject cwr_type;
2790
Tal Einatc4bccd32018-09-12 00:49:13 +03002791/*[clinic input]
2792@classmethod
2793itertools.combinations_with_replacement.__new__
2794 iterable: object
2795 r: Py_ssize_t
2796Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats.
2797
2798combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2799[clinic start generated code]*/
2800
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002801static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002802itertools_combinations_with_replacement_impl(PyTypeObject *type,
2803 PyObject *iterable,
2804 Py_ssize_t r)
2805/*[clinic end generated code: output=48b26856d4e659ca input=dc2a8c7ba785fad7]*/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002806{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002807 cwrobject *co;
2808 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002809 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002810 Py_ssize_t *indices = NULL;
2811 Py_ssize_t i;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002812
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002813 pool = PySequence_Tuple(iterable);
2814 if (pool == NULL)
2815 goto error;
2816 n = PyTuple_GET_SIZE(pool);
2817 if (r < 0) {
2818 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2819 goto error;
2820 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002821
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002822 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002823 if (indices == NULL) {
2824 PyErr_NoMemory();
2825 goto error;
2826 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002827
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002828 for (i=0 ; i<r ; i++)
2829 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002830
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002831 /* create cwrobject structure */
2832 co = (cwrobject *)type->tp_alloc(type, 0);
2833 if (co == NULL)
2834 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002835
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002836 co->pool = pool;
2837 co->indices = indices;
2838 co->result = NULL;
2839 co->r = r;
2840 co->stopped = !n && r;
2841
2842 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002843
2844error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002845 if (indices != NULL)
2846 PyMem_Free(indices);
2847 Py_XDECREF(pool);
2848 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002849}
2850
2851static void
2852cwr_dealloc(cwrobject *co)
2853{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002854 PyObject_GC_UnTrack(co);
2855 Py_XDECREF(co->pool);
2856 Py_XDECREF(co->result);
2857 if (co->indices != NULL)
2858 PyMem_Free(co->indices);
2859 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002860}
2861
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002862static PyObject *
2863cwr_sizeof(cwrobject *co, void *unused)
2864{
2865 Py_ssize_t res;
2866
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002867 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002868 res += co->r * sizeof(Py_ssize_t);
2869 return PyLong_FromSsize_t(res);
2870}
2871
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002872static int
2873cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2874{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002875 Py_VISIT(co->pool);
2876 Py_VISIT(co->result);
2877 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002878}
2879
2880static PyObject *
2881cwr_next(cwrobject *co)
2882{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002883 PyObject *elem;
2884 PyObject *oldelem;
2885 PyObject *pool = co->pool;
2886 Py_ssize_t *indices = co->indices;
2887 PyObject *result = co->result;
2888 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2889 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002890 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002891
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002892 if (co->stopped)
2893 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002894
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002895 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002896 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002897 result = PyTuple_New(r);
2898 if (result == NULL)
2899 goto empty;
2900 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002901 if (n > 0) {
2902 elem = PyTuple_GET_ITEM(pool, 0);
2903 for (i=0; i<r ; i++) {
2904 assert(indices[i] == 0);
2905 Py_INCREF(elem);
2906 PyTuple_SET_ITEM(result, i, elem);
2907 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002908 }
2909 } else {
2910 /* Copy the previous result tuple or re-use it if available */
2911 if (Py_REFCNT(result) > 1) {
2912 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002913 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002914 if (result == NULL)
2915 goto empty;
2916 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002917 Py_DECREF(old_result);
2918 }
2919 /* Now, we've got the only copy so we can update it in-place CPython's
2920 empty tuple is a singleton and cached in PyTuple's freelist. */
2921 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002922
Tim Peters9edb1682013-09-03 11:49:31 -05002923 /* Scan indices right-to-left until finding one that is not
2924 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002925 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2926 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002927
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002928 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002929 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002930 if (i < 0)
2931 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002932
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002933 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002934 maximum. Then set all to the right to the same value. */
2935 index = indices[i] + 1;
2936 assert(index < n);
2937 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002938 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002939 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002940 Py_INCREF(elem);
2941 oldelem = PyTuple_GET_ITEM(result, i);
2942 PyTuple_SET_ITEM(result, i, elem);
2943 Py_DECREF(oldelem);
2944 }
2945 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002946
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002947 Py_INCREF(result);
2948 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002949
2950empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002951 co->stopped = 1;
2952 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002953}
2954
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002955static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302956cwr_reduce(cwrobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002957{
2958 if (lz->result == NULL) {
2959 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2960 } else if (lz->stopped) {
2961 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2962 } else {
2963 PyObject *indices;
2964 Py_ssize_t i;
2965
2966 /* we must pickle the indices and use them for setstate */
2967 indices = PyTuple_New(lz->r);
2968 if (!indices)
2969 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002970 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002971 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2972 if (!index) {
2973 Py_DECREF(indices);
2974 return NULL;
2975 }
2976 PyTuple_SET_ITEM(indices, i, index);
2977 }
2978
2979 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2980 }
2981}
2982
2983static PyObject *
2984cwr_setstate(cwrobject *lz, PyObject *state)
2985{
2986 PyObject *result;
2987 Py_ssize_t n, i;
2988
2989 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2990 {
2991 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2992 return NULL;
2993 }
2994
2995 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002996 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002997 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2998 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002999
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003000 if (index < 0 && PyErr_Occurred())
3001 return NULL; /* not an integer */
3002 /* clamp the index */
3003 if (index < 0)
3004 index = 0;
3005 else if (index > n-1)
3006 index = n-1;
3007 lz->indices[i] = index;
3008 }
3009 result = PyTuple_New(lz->r);
3010 if (result == NULL)
3011 return NULL;
3012 for (i=0; i<lz->r; i++) {
3013 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
3014 Py_INCREF(element);
3015 PyTuple_SET_ITEM(result, i, element);
3016 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003017 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003018 Py_RETURN_NONE;
3019}
3020
3021static PyMethodDef cwr_methods[] = {
3022 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
3023 reduce_doc},
3024 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
3025 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003026 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
3027 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003028 {NULL, NULL} /* sentinel */
3029};
3030
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003031static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003032 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003033 "itertools.combinations_with_replacement", /* tp_name */
3034 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003035 0, /* tp_itemsize */
3036 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003037 (destructor)cwr_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003038 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003039 0, /* tp_getattr */
3040 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003041 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003042 0, /* tp_repr */
3043 0, /* tp_as_number */
3044 0, /* tp_as_sequence */
3045 0, /* tp_as_mapping */
3046 0, /* tp_hash */
3047 0, /* tp_call */
3048 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003049 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003050 0, /* tp_setattro */
3051 0, /* tp_as_buffer */
3052 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003053 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003054 itertools_combinations_with_replacement__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003055 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003056 0, /* tp_clear */
3057 0, /* tp_richcompare */
3058 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003059 PyObject_SelfIter, /* tp_iter */
3060 (iternextfunc)cwr_next, /* tp_iternext */
3061 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003062 0, /* tp_members */
3063 0, /* tp_getset */
3064 0, /* tp_base */
3065 0, /* tp_dict */
3066 0, /* tp_descr_get */
3067 0, /* tp_descr_set */
3068 0, /* tp_dictoffset */
3069 0, /* tp_init */
3070 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003071 itertools_combinations_with_replacement, /* tp_new */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003072 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003073};
3074
3075
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003076/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003077
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003078def permutations(iterable, r=None):
3079 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
3080 pool = tuple(iterable)
3081 n = len(pool)
3082 r = n if r is None else r
3083 indices = range(n)
3084 cycles = range(n-r+1, n+1)[::-1]
3085 yield tuple(pool[i] for i in indices[:r])
3086 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003087 for i in reversed(range(r)):
3088 cycles[i] -= 1
3089 if cycles[i] == 0:
3090 indices[i:] = indices[i+1:] + indices[i:i+1]
3091 cycles[i] = n - i
3092 else:
3093 j = cycles[i]
3094 indices[i], indices[-j] = indices[-j], indices[i]
3095 yield tuple(pool[i] for i in indices[:r])
3096 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003097 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003098 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003099*/
3100
3101typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003102 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003103 PyObject *pool; /* input converted to a tuple */
3104 Py_ssize_t *indices; /* one index per element in the pool */
3105 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3106 PyObject *result; /* most recently returned result tuple */
3107 Py_ssize_t r; /* size of result tuple */
3108 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003109} permutationsobject;
3110
3111static PyTypeObject permutations_type;
3112
Tal Einatc4bccd32018-09-12 00:49:13 +03003113/*[clinic input]
3114@classmethod
3115itertools.permutations.__new__
3116 iterable: object
3117 r as robj: object = None
3118Return successive r-length permutations of elements in the iterable.
3119
3120permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)
3121[clinic start generated code]*/
3122
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003123static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003124itertools_permutations_impl(PyTypeObject *type, PyObject *iterable,
3125 PyObject *robj)
3126/*[clinic end generated code: output=296a72fa76d620ea input=57d0170a4ac0ec7a]*/
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003127{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003128 permutationsobject *po;
3129 Py_ssize_t n;
3130 Py_ssize_t r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003131 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003132 Py_ssize_t *indices = NULL;
3133 Py_ssize_t *cycles = NULL;
3134 Py_ssize_t i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003135
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003136 pool = PySequence_Tuple(iterable);
3137 if (pool == NULL)
3138 goto error;
3139 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003140
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003141 r = n;
3142 if (robj != Py_None) {
3143 if (!PyLong_Check(robj)) {
3144 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3145 goto error;
3146 }
3147 r = PyLong_AsSsize_t(robj);
3148 if (r == -1 && PyErr_Occurred())
3149 goto error;
3150 }
3151 if (r < 0) {
3152 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3153 goto error;
3154 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003155
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003156 indices = PyMem_New(Py_ssize_t, n);
3157 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003158 if (indices == NULL || cycles == NULL) {
3159 PyErr_NoMemory();
3160 goto error;
3161 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003162
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003163 for (i=0 ; i<n ; i++)
3164 indices[i] = i;
3165 for (i=0 ; i<r ; i++)
3166 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003167
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003168 /* create permutationsobject structure */
3169 po = (permutationsobject *)type->tp_alloc(type, 0);
3170 if (po == NULL)
3171 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003172
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003173 po->pool = pool;
3174 po->indices = indices;
3175 po->cycles = cycles;
3176 po->result = NULL;
3177 po->r = r;
3178 po->stopped = r > n ? 1 : 0;
3179
3180 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003181
3182error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003183 if (indices != NULL)
3184 PyMem_Free(indices);
3185 if (cycles != NULL)
3186 PyMem_Free(cycles);
3187 Py_XDECREF(pool);
3188 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003189}
3190
3191static void
3192permutations_dealloc(permutationsobject *po)
3193{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003194 PyObject_GC_UnTrack(po);
3195 Py_XDECREF(po->pool);
3196 Py_XDECREF(po->result);
3197 PyMem_Free(po->indices);
3198 PyMem_Free(po->cycles);
3199 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003200}
3201
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003202static PyObject *
3203permutations_sizeof(permutationsobject *po, void *unused)
3204{
3205 Py_ssize_t res;
3206
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003207 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003208 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3209 res += po->r * sizeof(Py_ssize_t);
3210 return PyLong_FromSsize_t(res);
3211}
3212
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003213static int
3214permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3215{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003216 Py_VISIT(po->pool);
3217 Py_VISIT(po->result);
3218 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003219}
3220
3221static PyObject *
3222permutations_next(permutationsobject *po)
3223{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003224 PyObject *elem;
3225 PyObject *oldelem;
3226 PyObject *pool = po->pool;
3227 Py_ssize_t *indices = po->indices;
3228 Py_ssize_t *cycles = po->cycles;
3229 PyObject *result = po->result;
3230 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3231 Py_ssize_t r = po->r;
3232 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003233
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003234 if (po->stopped)
3235 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003236
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003237 if (result == NULL) {
3238 /* On the first pass, initialize result tuple using the indices */
3239 result = PyTuple_New(r);
3240 if (result == NULL)
3241 goto empty;
3242 po->result = result;
3243 for (i=0; i<r ; i++) {
3244 index = indices[i];
3245 elem = PyTuple_GET_ITEM(pool, index);
3246 Py_INCREF(elem);
3247 PyTuple_SET_ITEM(result, i, elem);
3248 }
3249 } else {
3250 if (n == 0)
3251 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003252
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003253 /* Copy the previous result tuple or re-use it if available */
3254 if (Py_REFCNT(result) > 1) {
3255 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05003256 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003257 if (result == NULL)
3258 goto empty;
3259 po->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003260 Py_DECREF(old_result);
3261 }
3262 /* Now, we've got the only copy so we can update it in-place */
3263 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003264
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003265 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3266 for (i=r-1 ; i>=0 ; i--) {
3267 cycles[i] -= 1;
3268 if (cycles[i] == 0) {
3269 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3270 index = indices[i];
3271 for (j=i ; j<n-1 ; j++)
3272 indices[j] = indices[j+1];
3273 indices[n-1] = index;
3274 cycles[i] = n - i;
3275 } else {
3276 j = cycles[i];
3277 index = indices[i];
3278 indices[i] = indices[n-j];
3279 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003280
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003281 for (k=i; k<r ; k++) {
3282 /* start with i, the leftmost element that changed */
3283 /* yield tuple(pool[k] for k in indices[:r]) */
3284 index = indices[k];
3285 elem = PyTuple_GET_ITEM(pool, index);
3286 Py_INCREF(elem);
3287 oldelem = PyTuple_GET_ITEM(result, k);
3288 PyTuple_SET_ITEM(result, k, elem);
3289 Py_DECREF(oldelem);
3290 }
3291 break;
3292 }
3293 }
3294 /* If i is negative, then the cycles have all
3295 rolled-over and we're done. */
3296 if (i < 0)
3297 goto empty;
3298 }
3299 Py_INCREF(result);
3300 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003301
3302empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003303 po->stopped = 1;
3304 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003305}
3306
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003307static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303308permutations_reduce(permutationsobject *po, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003309{
3310 if (po->result == NULL) {
3311 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3312 } else if (po->stopped) {
3313 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3314 } else {
3315 PyObject *indices=NULL, *cycles=NULL;
3316 Py_ssize_t n, i;
3317
3318 /* we must pickle the indices and cycles and use them for setstate */
3319 n = PyTuple_GET_SIZE(po->pool);
3320 indices = PyTuple_New(n);
3321 if (indices == NULL)
3322 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003323 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003324 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3325 if (!index)
3326 goto err;
3327 PyTuple_SET_ITEM(indices, i, index);
3328 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003329
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003330 cycles = PyTuple_New(po->r);
3331 if (cycles == NULL)
3332 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003333 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003334 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3335 if (!index)
3336 goto err;
3337 PyTuple_SET_ITEM(cycles, i, index);
3338 }
3339 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3340 po->pool, po->r,
3341 indices, cycles);
3342 err:
3343 Py_XDECREF(indices);
3344 Py_XDECREF(cycles);
3345 return NULL;
3346 }
3347}
3348
3349static PyObject *
3350permutations_setstate(permutationsobject *po, PyObject *state)
3351{
3352 PyObject *indices, *cycles, *result;
3353 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003354
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003355 if (!PyTuple_Check(state)) {
3356 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3357 return NULL;
3358 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003359 if (!PyArg_ParseTuple(state, "O!O!",
3360 &PyTuple_Type, &indices,
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003361 &PyTuple_Type, &cycles)) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003362 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003363 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003364
3365 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003366 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003367 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3368 return NULL;
3369 }
3370
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003371 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003372 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3373 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3374 if (index < 0 && PyErr_Occurred())
3375 return NULL; /* not an integer */
3376 /* clamp the index */
3377 if (index < 0)
3378 index = 0;
3379 else if (index > n-1)
3380 index = n-1;
3381 po->indices[i] = index;
3382 }
3383
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003384 for (i=0; i<po->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003385 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3386 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3387 if (index < 0 && PyErr_Occurred())
3388 return NULL; /* not an integer */
3389 if (index < 1)
3390 index = 1;
3391 else if (index > n-i)
3392 index = n-i;
3393 po->cycles[i] = index;
3394 }
3395 result = PyTuple_New(po->r);
3396 if (result == NULL)
3397 return NULL;
3398 for (i=0; i<po->r; i++) {
3399 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3400 Py_INCREF(element);
3401 PyTuple_SET_ITEM(result, i, element);
3402 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003403 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003404 Py_RETURN_NONE;
3405}
3406
3407static PyMethodDef permuations_methods[] = {
3408 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3409 reduce_doc},
3410 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3411 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003412 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3413 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003414 {NULL, NULL} /* sentinel */
3415};
3416
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003417static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003418 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003419 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003420 sizeof(permutationsobject), /* tp_basicsize */
3421 0, /* tp_itemsize */
3422 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003423 (destructor)permutations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003424 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003425 0, /* tp_getattr */
3426 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003427 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003428 0, /* tp_repr */
3429 0, /* tp_as_number */
3430 0, /* tp_as_sequence */
3431 0, /* tp_as_mapping */
3432 0, /* tp_hash */
3433 0, /* tp_call */
3434 0, /* tp_str */
3435 PyObject_GenericGetAttr, /* tp_getattro */
3436 0, /* tp_setattro */
3437 0, /* tp_as_buffer */
3438 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3439 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003440 itertools_permutations__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003441 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003442 0, /* tp_clear */
3443 0, /* tp_richcompare */
3444 0, /* tp_weaklistoffset */
3445 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003446 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003447 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003448 0, /* tp_members */
3449 0, /* tp_getset */
3450 0, /* tp_base */
3451 0, /* tp_dict */
3452 0, /* tp_descr_get */
3453 0, /* tp_descr_set */
3454 0, /* tp_dictoffset */
3455 0, /* tp_init */
3456 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003457 itertools_permutations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003458 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003459};
3460
Tal Einatc4bccd32018-09-12 00:49:13 +03003461
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003462/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003463
3464typedef struct {
3465 PyObject_HEAD
3466 PyObject *total;
3467 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003468 PyObject *binop;
Lisa Roach9718b592018-09-23 17:34:59 -07003469 PyObject *initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003470} accumulateobject;
3471
3472static PyTypeObject accumulate_type;
3473
Tal Einatc4bccd32018-09-12 00:49:13 +03003474/*[clinic input]
3475@classmethod
3476itertools.accumulate.__new__
3477 iterable: object
3478 func as binop: object = None
Lisa Roach9718b592018-09-23 17:34:59 -07003479 *
3480 initial: object = None
Tal Einatc4bccd32018-09-12 00:49:13 +03003481Return series of accumulated sums (or other binary function results).
3482[clinic start generated code]*/
3483
Raymond Hettinger482ba772010-12-01 22:48:00 +00003484static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003485itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
Lisa Roach9718b592018-09-23 17:34:59 -07003486 PyObject *binop, PyObject *initial)
3487/*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003488{
Raymond Hettinger482ba772010-12-01 22:48:00 +00003489 PyObject *it;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003490 accumulateobject *lz;
3491
Raymond Hettinger482ba772010-12-01 22:48:00 +00003492 /* Get iterator. */
3493 it = PyObject_GetIter(iterable);
3494 if (it == NULL)
3495 return NULL;
3496
Raymond Hettinger482ba772010-12-01 22:48:00 +00003497 /* create accumulateobject structure */
3498 lz = (accumulateobject *)type->tp_alloc(type, 0);
3499 if (lz == NULL) {
3500 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003501 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003502 }
3503
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003504 if (binop != Py_None) {
3505 Py_XINCREF(binop);
3506 lz->binop = binop;
3507 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003508 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003509 lz->it = it;
Lisa Roach9718b592018-09-23 17:34:59 -07003510 Py_XINCREF(initial);
3511 lz->initial = initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003512 return (PyObject *)lz;
3513}
3514
3515static void
3516accumulate_dealloc(accumulateobject *lz)
3517{
3518 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003519 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003520 Py_XDECREF(lz->total);
3521 Py_XDECREF(lz->it);
Lisa Roach9718b592018-09-23 17:34:59 -07003522 Py_XDECREF(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003523 Py_TYPE(lz)->tp_free(lz);
3524}
3525
3526static int
3527accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3528{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003529 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003530 Py_VISIT(lz->it);
3531 Py_VISIT(lz->total);
Lisa Roach9718b592018-09-23 17:34:59 -07003532 Py_VISIT(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003533 return 0;
3534}
3535
3536static PyObject *
3537accumulate_next(accumulateobject *lz)
3538{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003539 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003540
Lisa Roach9718b592018-09-23 17:34:59 -07003541 if (lz->initial != Py_None) {
3542 lz->total = lz->initial;
3543 Py_INCREF(Py_None);
3544 lz->initial = Py_None;
3545 Py_INCREF(lz->total);
3546 return lz->total;
3547 }
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003548 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003549 if (val == NULL)
3550 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003551
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003552 if (lz->total == NULL) {
3553 Py_INCREF(val);
3554 lz->total = val;
3555 return lz->total;
3556 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003557
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003558 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003559 newtotal = PyNumber_Add(lz->total, val);
3560 else
3561 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003562 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003563 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003564 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003565
Raymond Hettinger482ba772010-12-01 22:48:00 +00003566 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003567 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003568 return newtotal;
3569}
3570
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003571static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303572accumulate_reduce(accumulateobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003573{
Lisa Roach9718b592018-09-23 17:34:59 -07003574 if (lz->initial != Py_None) {
3575 PyObject *it;
3576
3577 assert(lz->total == NULL);
3578 if (PyType_Ready(&chain_type) < 0)
3579 return NULL;
3580 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3581 lz->initial, lz->it);
3582 if (it == NULL)
3583 return NULL;
3584 return Py_BuildValue("O(NO)O", Py_TYPE(lz),
3585 it, lz->binop?lz->binop:Py_None, Py_None);
3586 }
Serhiy Storchakad5516252016-03-06 14:00:45 +02003587 if (lz->total == Py_None) {
3588 PyObject *it;
3589
3590 if (PyType_Ready(&chain_type) < 0)
3591 return NULL;
3592 if (PyType_Ready(&islice_type) < 0)
3593 return NULL;
3594 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3595 lz->total, lz->it);
3596 if (it == NULL)
3597 return NULL;
3598 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3599 it, lz->binop ? lz->binop : Py_None);
3600 if (it == NULL)
3601 return NULL;
3602 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3603 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003604 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3605 lz->it, lz->binop?lz->binop:Py_None,
3606 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003607}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003608
3609static PyObject *
3610accumulate_setstate(accumulateobject *lz, PyObject *state)
3611{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003612 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003613 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003614 Py_RETURN_NONE;
3615}
3616
3617static PyMethodDef accumulate_methods[] = {
3618 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3619 reduce_doc},
3620 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3621 setstate_doc},
3622 {NULL, NULL} /* sentinel */
3623};
3624
Raymond Hettinger482ba772010-12-01 22:48:00 +00003625static PyTypeObject accumulate_type = {
3626 PyVarObject_HEAD_INIT(NULL, 0)
3627 "itertools.accumulate", /* tp_name */
3628 sizeof(accumulateobject), /* tp_basicsize */
3629 0, /* tp_itemsize */
3630 /* methods */
3631 (destructor)accumulate_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003632 0, /* tp_vectorcall_offset */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003633 0, /* tp_getattr */
3634 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003635 0, /* tp_as_async */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003636 0, /* tp_repr */
3637 0, /* tp_as_number */
3638 0, /* tp_as_sequence */
3639 0, /* tp_as_mapping */
3640 0, /* tp_hash */
3641 0, /* tp_call */
3642 0, /* tp_str */
3643 PyObject_GenericGetAttr, /* tp_getattro */
3644 0, /* tp_setattro */
3645 0, /* tp_as_buffer */
3646 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3647 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003648 itertools_accumulate__doc__, /* tp_doc */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003649 (traverseproc)accumulate_traverse, /* tp_traverse */
3650 0, /* tp_clear */
3651 0, /* tp_richcompare */
3652 0, /* tp_weaklistoffset */
3653 PyObject_SelfIter, /* tp_iter */
3654 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003655 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003656 0, /* tp_members */
3657 0, /* tp_getset */
3658 0, /* tp_base */
3659 0, /* tp_dict */
3660 0, /* tp_descr_get */
3661 0, /* tp_descr_set */
3662 0, /* tp_dictoffset */
3663 0, /* tp_init */
3664 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003665 itertools_accumulate, /* tp_new */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003666 PyObject_GC_Del, /* tp_free */
3667};
3668
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003669
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003670/* compress object ************************************************************/
3671
3672/* Equivalent to:
3673
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003674 def compress(data, selectors):
3675 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3676 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003677*/
3678
3679typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003680 PyObject_HEAD
3681 PyObject *data;
3682 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003683} compressobject;
3684
3685static PyTypeObject compress_type;
3686
Tal Einatc4bccd32018-09-12 00:49:13 +03003687/*[clinic input]
3688@classmethod
3689itertools.compress.__new__
3690 data as seq1: object
3691 selectors as seq2: object
3692Return data elements corresponding to true selector elements.
3693
3694Forms a shorter iterator from selected data elements using the selectors to
3695choose the data elements.
3696[clinic start generated code]*/
3697
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003698static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003699itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
3700/*[clinic end generated code: output=7e67157212ed09e0 input=79596d7cd20c77e5]*/
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003701{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003702 PyObject *data=NULL, *selectors=NULL;
3703 compressobject *lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003704
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003705 data = PyObject_GetIter(seq1);
3706 if (data == NULL)
3707 goto fail;
3708 selectors = PyObject_GetIter(seq2);
3709 if (selectors == NULL)
3710 goto fail;
3711
3712 /* create compressobject structure */
3713 lz = (compressobject *)type->tp_alloc(type, 0);
3714 if (lz == NULL)
3715 goto fail;
3716 lz->data = data;
3717 lz->selectors = selectors;
3718 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003719
3720fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003721 Py_XDECREF(data);
3722 Py_XDECREF(selectors);
3723 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003724}
3725
3726static void
3727compress_dealloc(compressobject *lz)
3728{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003729 PyObject_GC_UnTrack(lz);
3730 Py_XDECREF(lz->data);
3731 Py_XDECREF(lz->selectors);
3732 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003733}
3734
3735static int
3736compress_traverse(compressobject *lz, visitproc visit, void *arg)
3737{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003738 Py_VISIT(lz->data);
3739 Py_VISIT(lz->selectors);
3740 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003741}
3742
3743static PyObject *
3744compress_next(compressobject *lz)
3745{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003746 PyObject *data = lz->data, *selectors = lz->selectors;
3747 PyObject *datum, *selector;
3748 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3749 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3750 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003751
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003752 while (1) {
3753 /* Steps: get datum, get selector, evaluate selector.
3754 Order is important (to match the pure python version
3755 in terms of which input gets a chance to raise an
3756 exception first).
3757 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003758
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003759 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003760 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003761 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003762
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003763 selector = selectornext(selectors);
3764 if (selector == NULL) {
3765 Py_DECREF(datum);
3766 return NULL;
3767 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003768
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003769 ok = PyObject_IsTrue(selector);
3770 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003771 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003772 return datum;
3773 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003774 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003775 return NULL;
3776 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003777}
3778
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003779static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303780compress_reduce(compressobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003781{
3782 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3783 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003784}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003785
3786static PyMethodDef compress_methods[] = {
3787 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3788 reduce_doc},
3789 {NULL, NULL} /* sentinel */
3790};
3791
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003792static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003793 PyVarObject_HEAD_INIT(NULL, 0)
3794 "itertools.compress", /* tp_name */
3795 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003796 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003797 /* methods */
3798 (destructor)compress_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003799 0, /* tp_vectorcall_offset */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003800 0, /* tp_getattr */
3801 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003802 0, /* tp_as_async */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003803 0, /* tp_repr */
3804 0, /* tp_as_number */
3805 0, /* tp_as_sequence */
3806 0, /* tp_as_mapping */
3807 0, /* tp_hash */
3808 0, /* tp_call */
3809 0, /* tp_str */
3810 PyObject_GenericGetAttr, /* tp_getattro */
3811 0, /* tp_setattro */
3812 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003813 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003814 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003815 itertools_compress__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003816 (traverseproc)compress_traverse, /* tp_traverse */
3817 0, /* tp_clear */
3818 0, /* tp_richcompare */
3819 0, /* tp_weaklistoffset */
3820 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003821 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003822 compress_methods, /* tp_methods */
3823 0, /* tp_members */
3824 0, /* tp_getset */
3825 0, /* tp_base */
3826 0, /* tp_dict */
3827 0, /* tp_descr_get */
3828 0, /* tp_descr_set */
3829 0, /* tp_dictoffset */
3830 0, /* tp_init */
3831 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003832 itertools_compress, /* tp_new */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003833 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003834};
3835
3836
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003837/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003838
3839typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003840 PyObject_HEAD
3841 PyObject *func;
3842 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003843} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003844
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003845static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003846
Tal Einatc4bccd32018-09-12 00:49:13 +03003847/*[clinic input]
3848@classmethod
3849itertools.filterfalse.__new__
3850 function as func: object
3851 iterable as seq: object
3852 /
3853Return those items of iterable for which function(item) is false.
3854
3855If function is None, return the items that are false.
3856[clinic start generated code]*/
3857
Raymond Hettinger60eca932003-02-09 06:40:58 +00003858static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003859itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
3860/*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003861{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003862 PyObject *it;
3863 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003864
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003865 /* Get iterator. */
3866 it = PyObject_GetIter(seq);
3867 if (it == NULL)
3868 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003869
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003870 /* create filterfalseobject structure */
3871 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3872 if (lz == NULL) {
3873 Py_DECREF(it);
3874 return NULL;
3875 }
3876 Py_INCREF(func);
3877 lz->func = func;
3878 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003879
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003880 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003881}
3882
3883static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003884filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003885{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003886 PyObject_GC_UnTrack(lz);
3887 Py_XDECREF(lz->func);
3888 Py_XDECREF(lz->it);
3889 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003890}
3891
3892static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003893filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003894{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003895 Py_VISIT(lz->it);
3896 Py_VISIT(lz->func);
3897 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003898}
3899
3900static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003901filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003902{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003903 PyObject *item;
3904 PyObject *it = lz->it;
3905 long ok;
3906 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003907
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003908 iternext = *Py_TYPE(it)->tp_iternext;
3909 for (;;) {
3910 item = iternext(it);
3911 if (item == NULL)
3912 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003913
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003914 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3915 ok = PyObject_IsTrue(item);
3916 } else {
3917 PyObject *good;
Jeroen Demeyer196a5302019-07-04 12:31:34 +02003918 good = _PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003919 if (good == NULL) {
3920 Py_DECREF(item);
3921 return NULL;
3922 }
3923 ok = PyObject_IsTrue(good);
3924 Py_DECREF(good);
3925 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003926 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003927 return item;
3928 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003929 if (ok < 0)
3930 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003931 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003932}
3933
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003934static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303935filterfalse_reduce(filterfalseobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003936{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003937 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3938}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003939
3940static PyMethodDef filterfalse_methods[] = {
3941 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3942 reduce_doc},
3943 {NULL, NULL} /* sentinel */
3944};
3945
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003946static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003947 PyVarObject_HEAD_INIT(NULL, 0)
3948 "itertools.filterfalse", /* tp_name */
3949 sizeof(filterfalseobject), /* tp_basicsize */
3950 0, /* tp_itemsize */
3951 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003952 (destructor)filterfalse_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003953 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003954 0, /* tp_getattr */
3955 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003956 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003957 0, /* tp_repr */
3958 0, /* tp_as_number */
3959 0, /* tp_as_sequence */
3960 0, /* tp_as_mapping */
3961 0, /* tp_hash */
3962 0, /* tp_call */
3963 0, /* tp_str */
3964 PyObject_GenericGetAttr, /* tp_getattro */
3965 0, /* tp_setattro */
3966 0, /* tp_as_buffer */
3967 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3968 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003969 itertools_filterfalse__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003970 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003971 0, /* tp_clear */
3972 0, /* tp_richcompare */
3973 0, /* tp_weaklistoffset */
3974 PyObject_SelfIter, /* tp_iter */
3975 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003976 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003977 0, /* tp_members */
3978 0, /* tp_getset */
3979 0, /* tp_base */
3980 0, /* tp_dict */
3981 0, /* tp_descr_get */
3982 0, /* tp_descr_set */
3983 0, /* tp_dictoffset */
3984 0, /* tp_init */
3985 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003986 itertools_filterfalse, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003987 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003988};
3989
3990
3991/* count object ************************************************************/
3992
3993typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003994 PyObject_HEAD
3995 Py_ssize_t cnt;
3996 PyObject *long_cnt;
3997 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003998} countobject;
3999
Raymond Hettinger30729212009-02-12 06:28:27 +00004000/* Counting logic and invariants:
4001
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004002fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00004003
Serhiy Storchaka483405b2015-02-17 10:14:30 +02004004 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004005 Advances with: cnt += 1
4006 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00004007
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004008slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00004009
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004010 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
4011 All counting is done with python objects (no overflows or underflows).
4012 Advances with: long_cnt += long_step
4013 Step may be zero -- effectively a slow version of repeat(cnt).
4014 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00004015*/
4016
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004017static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004018
Tal Einatc4bccd32018-09-12 00:49:13 +03004019/*[clinic input]
4020@classmethod
4021itertools.count.__new__
4022 start as long_cnt: object(c_default="NULL") = 0
4023 step as long_step: object(c_default="NULL") = 1
4024Return a count object whose .__next__() method returns consecutive values.
4025
4026Equivalent to:
4027 def count(firstval=0, step=1):
4028 x = firstval
4029 while 1:
4030 yield x
4031 x += step
4032[clinic start generated code]*/
4033
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004034static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03004035itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
4036 PyObject *long_step)
4037/*[clinic end generated code: output=09a9250aebd00b1c input=d7a85eec18bfcd94]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004038{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004039 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004040 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004041 Py_ssize_t cnt = 0;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00004042 long step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004043
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004044 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
4045 (long_step != NULL && !PyNumber_Check(long_step))) {
4046 PyErr_SetString(PyExc_TypeError, "a number is required");
4047 return NULL;
4048 }
Raymond Hettinger30729212009-02-12 06:28:27 +00004049
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004050 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
4051 (long_step == NULL || PyLong_Check(long_step));
4052
4053 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004054 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004055 if (fast_mode) {
4056 assert(PyLong_Check(long_cnt));
4057 cnt = PyLong_AsSsize_t(long_cnt);
4058 if (cnt == -1 && PyErr_Occurred()) {
4059 PyErr_Clear();
4060 fast_mode = 0;
4061 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004062 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004063 } else {
4064 cnt = 0;
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004065 long_cnt = _PyLong_Zero;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004066 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004067 Py_INCREF(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004068
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004069 /* If not specified, step defaults to 1 */
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004070 if (long_step == NULL)
4071 long_step = _PyLong_One;
4072 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004073
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004074 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004075
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004076 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004077 if (fast_mode) {
4078 assert(PyLong_Check(long_step));
4079 step = PyLong_AsLong(long_step);
4080 if (step != 1) {
4081 fast_mode = 0;
4082 if (step == -1 && PyErr_Occurred())
4083 PyErr_Clear();
4084 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004085 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004086
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004087 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004088 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004089 else
4090 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004091
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004092 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4093 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4094 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004095 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004096
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004097 /* create countobject structure */
4098 lz = (countobject *)type->tp_alloc(type, 0);
4099 if (lz == NULL) {
4100 Py_XDECREF(long_cnt);
Zackery Spytz0523c392019-03-26 00:05:29 -06004101 Py_DECREF(long_step);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004102 return NULL;
4103 }
4104 lz->cnt = cnt;
4105 lz->long_cnt = long_cnt;
4106 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004107
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004108 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004109}
4110
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004111static void
4112count_dealloc(countobject *lz)
4113{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004114 PyObject_GC_UnTrack(lz);
4115 Py_XDECREF(lz->long_cnt);
4116 Py_XDECREF(lz->long_step);
4117 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004118}
4119
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004120static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004121count_traverse(countobject *lz, visitproc visit, void *arg)
4122{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004123 Py_VISIT(lz->long_cnt);
4124 Py_VISIT(lz->long_step);
4125 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004126}
4127
4128static PyObject *
4129count_nextlong(countobject *lz)
4130{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004131 PyObject *long_cnt;
4132 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004133
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004134 long_cnt = lz->long_cnt;
4135 if (long_cnt == NULL) {
4136 /* Switch to slow_mode */
4137 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4138 if (long_cnt == NULL)
4139 return NULL;
4140 }
4141 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004142
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004143 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4144 if (stepped_up == NULL)
4145 return NULL;
4146 lz->long_cnt = stepped_up;
4147 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004148}
4149
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004150static PyObject *
4151count_next(countobject *lz)
4152{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004153 if (lz->cnt == PY_SSIZE_T_MAX)
4154 return count_nextlong(lz);
4155 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004156}
4157
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004158static PyObject *
4159count_repr(countobject *lz)
4160{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004161 if (lz->cnt != PY_SSIZE_T_MAX)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004162 return PyUnicode_FromFormat("%s(%zd)",
4163 _PyType_Name(Py_TYPE(lz)), lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004164
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004165 if (PyLong_Check(lz->long_step)) {
4166 long step = PyLong_AsLong(lz->long_step);
4167 if (step == -1 && PyErr_Occurred()) {
4168 PyErr_Clear();
4169 }
4170 if (step == 1) {
4171 /* Don't display step when it is an integer equal to 1 */
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004172 return PyUnicode_FromFormat("%s(%R)",
4173 _PyType_Name(Py_TYPE(lz)),
4174 lz->long_cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004175 }
4176 }
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004177 return PyUnicode_FromFormat("%s(%R, %R)",
4178 _PyType_Name(Py_TYPE(lz)),
4179 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004180}
4181
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004182static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304183count_reduce(countobject *lz, PyObject *Py_UNUSED(ignored))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004184{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004185 if (lz->cnt == PY_SSIZE_T_MAX)
4186 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4187 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004188}
4189
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004190static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004191 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004192 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004193 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004194};
4195
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004196static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004197 PyVarObject_HEAD_INIT(NULL, 0)
4198 "itertools.count", /* tp_name */
4199 sizeof(countobject), /* tp_basicsize */
4200 0, /* tp_itemsize */
4201 /* methods */
4202 (destructor)count_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004203 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004204 0, /* tp_getattr */
4205 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004206 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004207 (reprfunc)count_repr, /* tp_repr */
4208 0, /* tp_as_number */
4209 0, /* tp_as_sequence */
4210 0, /* tp_as_mapping */
4211 0, /* tp_hash */
4212 0, /* tp_call */
4213 0, /* tp_str */
4214 PyObject_GenericGetAttr, /* tp_getattro */
4215 0, /* tp_setattro */
4216 0, /* tp_as_buffer */
4217 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004218 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03004219 itertools_count__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004220 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004221 0, /* tp_clear */
4222 0, /* tp_richcompare */
4223 0, /* tp_weaklistoffset */
4224 PyObject_SelfIter, /* tp_iter */
4225 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004226 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004227 0, /* tp_members */
4228 0, /* tp_getset */
4229 0, /* tp_base */
4230 0, /* tp_dict */
4231 0, /* tp_descr_get */
4232 0, /* tp_descr_set */
4233 0, /* tp_dictoffset */
4234 0, /* tp_init */
4235 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03004236 itertools_count, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004237 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004238};
4239
4240
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004241/* repeat object ************************************************************/
4242
4243typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004244 PyObject_HEAD
4245 PyObject *element;
4246 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004247} repeatobject;
4248
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004249static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004250
4251static PyObject *
4252repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4253{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004254 repeatobject *ro;
4255 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004256 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004257 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004258
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004259 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4260 &element, &cnt))
4261 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004262
Raymond Hettinger97d35552014-06-24 21:36:58 -07004263 if (kwds != NULL)
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004264 n_kwds = PyDict_GET_SIZE(kwds);
Raymond Hettinger97d35552014-06-24 21:36:58 -07004265 /* Does user supply times argument? */
4266 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004267 cnt = 0;
4268
4269 ro = (repeatobject *)type->tp_alloc(type, 0);
4270 if (ro == NULL)
4271 return NULL;
4272 Py_INCREF(element);
4273 ro->element = element;
4274 ro->cnt = cnt;
4275 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004276}
4277
4278static void
4279repeat_dealloc(repeatobject *ro)
4280{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004281 PyObject_GC_UnTrack(ro);
4282 Py_XDECREF(ro->element);
4283 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004284}
4285
4286static int
4287repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4288{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004289 Py_VISIT(ro->element);
4290 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004291}
4292
4293static PyObject *
4294repeat_next(repeatobject *ro)
4295{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004296 if (ro->cnt == 0)
4297 return NULL;
4298 if (ro->cnt > 0)
4299 ro->cnt--;
4300 Py_INCREF(ro->element);
4301 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004302}
4303
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004304static PyObject *
4305repeat_repr(repeatobject *ro)
4306{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004307 if (ro->cnt == -1)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004308 return PyUnicode_FromFormat("%s(%R)",
4309 _PyType_Name(Py_TYPE(ro)), ro->element);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004310 else
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004311 return PyUnicode_FromFormat("%s(%R, %zd)",
4312 _PyType_Name(Py_TYPE(ro)), ro->element,
4313 ro->cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004314}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004315
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004316static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304317repeat_len(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004318{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004319 if (ro->cnt == -1) {
4320 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4321 return NULL;
4322 }
4323 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004324}
4325
Armin Rigof5b3e362006-02-11 21:32:43 +00004326PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004327
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004328static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304329repeat_reduce(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004330{
4331 /* unpickle this so that a new repeat iterator is constructed with an
4332 * object, then call __setstate__ on it to set cnt
4333 */
4334 if (ro->cnt >= 0)
4335 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4336 else
4337 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4338}
4339
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004340static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004341 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004342 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004343 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004344};
4345
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004346PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004347"repeat(object [,times]) -> create an iterator which returns the object\n\
4348for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004349endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004350
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004351static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004352 PyVarObject_HEAD_INIT(NULL, 0)
4353 "itertools.repeat", /* tp_name */
4354 sizeof(repeatobject), /* tp_basicsize */
4355 0, /* tp_itemsize */
4356 /* methods */
4357 (destructor)repeat_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004358 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004359 0, /* tp_getattr */
4360 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004361 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004362 (reprfunc)repeat_repr, /* tp_repr */
4363 0, /* tp_as_number */
4364 0, /* tp_as_sequence */
4365 0, /* tp_as_mapping */
4366 0, /* tp_hash */
4367 0, /* tp_call */
4368 0, /* tp_str */
4369 PyObject_GenericGetAttr, /* tp_getattro */
4370 0, /* tp_setattro */
4371 0, /* tp_as_buffer */
4372 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4373 Py_TPFLAGS_BASETYPE, /* tp_flags */
4374 repeat_doc, /* tp_doc */
4375 (traverseproc)repeat_traverse, /* tp_traverse */
4376 0, /* tp_clear */
4377 0, /* tp_richcompare */
4378 0, /* tp_weaklistoffset */
4379 PyObject_SelfIter, /* tp_iter */
4380 (iternextfunc)repeat_next, /* tp_iternext */
4381 repeat_methods, /* tp_methods */
4382 0, /* tp_members */
4383 0, /* tp_getset */
4384 0, /* tp_base */
4385 0, /* tp_dict */
4386 0, /* tp_descr_get */
4387 0, /* tp_descr_set */
4388 0, /* tp_dictoffset */
4389 0, /* tp_init */
4390 0, /* tp_alloc */
4391 repeat_new, /* tp_new */
4392 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004393};
4394
Tal Einatc4bccd32018-09-12 00:49:13 +03004395
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004396/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004397
4398typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004399 PyObject_HEAD
4400 Py_ssize_t tuplesize;
4401 Py_ssize_t numactive;
4402 PyObject *ittuple; /* tuple of iterators */
4403 PyObject *result;
4404 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004405} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004406
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004407static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004408
4409static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004410zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004411{
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004412 _Py_IDENTIFIER(fillvalue);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004413 ziplongestobject *lz;
4414 Py_ssize_t i;
4415 PyObject *ittuple; /* tuple of iterators */
4416 PyObject *result;
4417 PyObject *fillvalue = Py_None;
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004418 Py_ssize_t tuplesize;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004419
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004420 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004421 fillvalue = NULL;
4422 if (PyDict_GET_SIZE(kwds) == 1) {
4423 fillvalue = _PyDict_GetItemIdWithError(kwds, &PyId_fillvalue);
4424 }
4425 if (fillvalue == NULL) {
4426 if (!PyErr_Occurred()) {
4427 PyErr_SetString(PyExc_TypeError,
4428 "zip_longest() got an unexpected keyword argument");
4429 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004430 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004431 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004432 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004433
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004434 /* args must be a tuple */
4435 assert(PyTuple_Check(args));
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004436 tuplesize = PyTuple_GET_SIZE(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004437
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004438 /* obtain iterators */
4439 ittuple = PyTuple_New(tuplesize);
4440 if (ittuple == NULL)
4441 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004442 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004443 PyObject *item = PyTuple_GET_ITEM(args, i);
4444 PyObject *it = PyObject_GetIter(item);
4445 if (it == NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004446 Py_DECREF(ittuple);
4447 return NULL;
4448 }
4449 PyTuple_SET_ITEM(ittuple, i, it);
4450 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004451
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004452 /* create a result holder */
4453 result = PyTuple_New(tuplesize);
4454 if (result == NULL) {
4455 Py_DECREF(ittuple);
4456 return NULL;
4457 }
4458 for (i=0 ; i < tuplesize ; i++) {
4459 Py_INCREF(Py_None);
4460 PyTuple_SET_ITEM(result, i, Py_None);
4461 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004462
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004463 /* create ziplongestobject structure */
4464 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4465 if (lz == NULL) {
4466 Py_DECREF(ittuple);
4467 Py_DECREF(result);
4468 return NULL;
4469 }
4470 lz->ittuple = ittuple;
4471 lz->tuplesize = tuplesize;
4472 lz->numactive = tuplesize;
4473 lz->result = result;
4474 Py_INCREF(fillvalue);
4475 lz->fillvalue = fillvalue;
4476 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004477}
4478
4479static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004480zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004481{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004482 PyObject_GC_UnTrack(lz);
4483 Py_XDECREF(lz->ittuple);
4484 Py_XDECREF(lz->result);
4485 Py_XDECREF(lz->fillvalue);
4486 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004487}
4488
4489static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004490zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004491{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004492 Py_VISIT(lz->ittuple);
4493 Py_VISIT(lz->result);
4494 Py_VISIT(lz->fillvalue);
4495 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004496}
4497
4498static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004499zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004500{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004501 Py_ssize_t i;
4502 Py_ssize_t tuplesize = lz->tuplesize;
4503 PyObject *result = lz->result;
4504 PyObject *it;
4505 PyObject *item;
4506 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004507
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004508 if (tuplesize == 0)
4509 return NULL;
4510 if (lz->numactive == 0)
4511 return NULL;
4512 if (Py_REFCNT(result) == 1) {
4513 Py_INCREF(result);
4514 for (i=0 ; i < tuplesize ; i++) {
4515 it = PyTuple_GET_ITEM(lz->ittuple, i);
4516 if (it == NULL) {
4517 Py_INCREF(lz->fillvalue);
4518 item = lz->fillvalue;
4519 } else {
4520 item = PyIter_Next(it);
4521 if (item == NULL) {
4522 lz->numactive -= 1;
4523 if (lz->numactive == 0 || PyErr_Occurred()) {
4524 lz->numactive = 0;
4525 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004526 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004527 } else {
4528 Py_INCREF(lz->fillvalue);
4529 item = lz->fillvalue;
4530 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4531 Py_DECREF(it);
4532 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004533 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004534 }
4535 olditem = PyTuple_GET_ITEM(result, i);
4536 PyTuple_SET_ITEM(result, i, item);
4537 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004538 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004539 } else {
4540 result = PyTuple_New(tuplesize);
4541 if (result == NULL)
4542 return NULL;
4543 for (i=0 ; i < tuplesize ; i++) {
4544 it = PyTuple_GET_ITEM(lz->ittuple, i);
4545 if (it == NULL) {
4546 Py_INCREF(lz->fillvalue);
4547 item = lz->fillvalue;
4548 } else {
4549 item = PyIter_Next(it);
4550 if (item == NULL) {
4551 lz->numactive -= 1;
4552 if (lz->numactive == 0 || PyErr_Occurred()) {
4553 lz->numactive = 0;
4554 Py_DECREF(result);
4555 return NULL;
4556 } else {
4557 Py_INCREF(lz->fillvalue);
4558 item = lz->fillvalue;
4559 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4560 Py_DECREF(it);
4561 }
4562 }
4563 }
4564 PyTuple_SET_ITEM(result, i, item);
4565 }
4566 }
4567 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004568}
4569
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004570static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304571zip_longest_reduce(ziplongestobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004572{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004573
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004574 /* Create a new tuple with empty sequences where appropriate to pickle.
4575 * Then use setstate to set the fillvalue
4576 */
4577 int i;
4578 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004579
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004580 if (args == NULL)
4581 return NULL;
4582 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4583 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4584 if (elem == NULL) {
4585 elem = PyTuple_New(0);
4586 if (elem == NULL) {
4587 Py_DECREF(args);
4588 return NULL;
4589 }
4590 } else
4591 Py_INCREF(elem);
4592 PyTuple_SET_ITEM(args, i, elem);
4593 }
4594 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4595}
4596
4597static PyObject *
4598zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4599{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004600 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004601 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004602 Py_RETURN_NONE;
4603}
4604
4605static PyMethodDef zip_longest_methods[] = {
4606 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4607 reduce_doc},
4608 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4609 setstate_doc},
4610 {NULL, NULL} /* sentinel */
4611};
4612
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004613PyDoc_STRVAR(zip_longest_doc,
4614"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004615\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004616Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004617the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004618method continues until the longest iterable in the argument sequence\n\
4619is exhausted and then it raises StopIteration. When the shorter iterables\n\
4620are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4621defaults to None or can be specified by a keyword argument.\n\
4622");
4623
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004624static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004625 PyVarObject_HEAD_INIT(NULL, 0)
4626 "itertools.zip_longest", /* tp_name */
4627 sizeof(ziplongestobject), /* tp_basicsize */
4628 0, /* tp_itemsize */
4629 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004630 (destructor)zip_longest_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004631 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004632 0, /* tp_getattr */
4633 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004634 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004635 0, /* tp_repr */
4636 0, /* tp_as_number */
4637 0, /* tp_as_sequence */
4638 0, /* tp_as_mapping */
4639 0, /* tp_hash */
4640 0, /* tp_call */
4641 0, /* tp_str */
4642 PyObject_GenericGetAttr, /* tp_getattro */
4643 0, /* tp_setattro */
4644 0, /* tp_as_buffer */
4645 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4646 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004647 zip_longest_doc, /* tp_doc */
4648 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004649 0, /* tp_clear */
4650 0, /* tp_richcompare */
4651 0, /* tp_weaklistoffset */
4652 PyObject_SelfIter, /* tp_iter */
4653 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004654 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004655 0, /* tp_members */
4656 0, /* tp_getset */
4657 0, /* tp_base */
4658 0, /* tp_dict */
4659 0, /* tp_descr_get */
4660 0, /* tp_descr_set */
4661 0, /* tp_dictoffset */
4662 0, /* tp_init */
4663 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004664 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004665 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004666};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004667
Tal Einatc4bccd32018-09-12 00:49:13 +03004668
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004669/* module level code ********************************************************/
4670
4671PyDoc_STRVAR(module_doc,
4672"Functional tools for creating and using iterators.\n\
4673\n\
4674Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004675count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004676cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004677repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004678\n\
4679Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004680accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
oldkaa0735f2018-02-02 16:52:55 +08004681chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
4682chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004683compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4684dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4685groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004686filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004687islice(seq, [start,] stop [, step]) --> elements from\n\
4688 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004689starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004690tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004691takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
oldkaa0735f2018-02-02 16:52:55 +08004692zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004693\n\
4694Combinatoric generators:\n\
4695product(p, q, ... [repeat=1]) --> cartesian product\n\
4696permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004697combinations(p, r)\n\
4698combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004699");
4700
4701
Raymond Hettingerad983e72003-11-12 14:32:26 +00004702static PyMethodDef module_methods[] = {
Tal Einatc4bccd32018-09-12 00:49:13 +03004703 ITERTOOLS_TEE_METHODDEF
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004704 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004705};
4706
Martin v. Löwis1a214512008-06-11 05:26:20 +00004707
4708static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004709 PyModuleDef_HEAD_INIT,
4710 "itertools",
4711 module_doc,
4712 -1,
4713 module_methods,
4714 NULL,
4715 NULL,
4716 NULL,
4717 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004718};
4719
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004720PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004721PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004722{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004723 int i;
4724 PyObject *m;
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004725 const char *name;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004726 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004727 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004728 &combinations_type,
4729 &cwr_type,
4730 &cycle_type,
4731 &dropwhile_type,
4732 &takewhile_type,
4733 &islice_type,
4734 &starmap_type,
4735 &chain_type,
4736 &compress_type,
4737 &filterfalse_type,
4738 &count_type,
4739 &ziplongest_type,
4740 &permutations_type,
4741 &product_type,
4742 &repeat_type,
4743 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004744 &_grouper_type,
4745 &tee_type,
4746 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004747 NULL
4748 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004749
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004750 Py_TYPE(&teedataobject_type) = &PyType_Type;
4751 m = PyModule_Create(&itertoolsmodule);
4752 if (m == NULL)
4753 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004754
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004755 for (i=0 ; typelist[i] != NULL ; i++) {
4756 if (PyType_Ready(typelist[i]) < 0)
4757 return NULL;
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004758 name = _PyType_Name(typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004759 Py_INCREF(typelist[i]);
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004760 PyModule_AddObject(m, name, (PyObject *)typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004761 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004762
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004763 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004764}