blob: 2e398425a1463af64613f31e4a84405eb10e075c [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 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000446 PyObject *nextlink;
447 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000448} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000449
450typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000451 PyObject_HEAD
452 teedataobject *dataobj;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700453 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000454 PyObject *weakreflist;
HongWeipengfa220ec2019-08-29 11:39:25 +0800455 unsigned long thread_id;
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
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000469 tdo->numread = 0;
470 tdo->nextlink = NULL;
471 Py_INCREF(it);
472 tdo->it = it;
473 PyObject_GC_Track(tdo);
474 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000475}
476
477static PyObject *
478teedataobject_jumplink(teedataobject *tdo)
479{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000480 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000481 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000482 Py_XINCREF(tdo->nextlink);
483 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000484}
485
486static PyObject *
487teedataobject_getitem(teedataobject *tdo, int i)
488{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000489 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000490
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000491 assert(i < LINKCELLS);
492 if (i < tdo->numread)
493 value = tdo->values[i];
494 else {
495 /* this is the lead iterator, so fetch more data */
496 assert(i == tdo->numread);
497 value = PyIter_Next(tdo->it);
498 if (value == NULL)
499 return NULL;
500 tdo->numread++;
501 tdo->values[i] = value;
502 }
503 Py_INCREF(value);
504 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000505}
506
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000507static int
508teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
509{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000510 int i;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700511
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000512 Py_VISIT(tdo->it);
513 for (i = 0; i < tdo->numread; i++)
514 Py_VISIT(tdo->values[i]);
515 Py_VISIT(tdo->nextlink);
516 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000517}
518
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200519static void
520teedataobject_safe_decref(PyObject *obj)
521{
522 while (obj && Py_TYPE(obj) == &teedataobject_type &&
523 Py_REFCNT(obj) == 1) {
524 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
525 ((teedataobject *)obj)->nextlink = NULL;
526 Py_DECREF(obj);
527 obj = nextlink;
528 }
529 Py_XDECREF(obj);
530}
531
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000532static int
533teedataobject_clear(teedataobject *tdo)
534{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000535 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200536 PyObject *tmp;
537
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000538 Py_CLEAR(tdo->it);
539 for (i=0 ; i<tdo->numread ; i++)
540 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200541 tmp = tdo->nextlink;
542 tdo->nextlink = NULL;
543 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000544 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000545}
546
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000547static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000548teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000549{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000550 PyObject_GC_UnTrack(tdo);
551 teedataobject_clear(tdo);
552 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000553}
554
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000555static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530556teedataobject_reduce(teedataobject *tdo, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000557{
558 int i;
559 /* create a temporary list of already iterated values */
560 PyObject *values = PyList_New(tdo->numread);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700561
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000562 if (!values)
563 return NULL;
564 for (i=0 ; i<tdo->numread ; i++) {
565 Py_INCREF(tdo->values[i]);
566 PyList_SET_ITEM(values, i, tdo->values[i]);
567 }
568 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
569 values,
570 tdo->nextlink ? tdo->nextlink : Py_None);
571}
572
Tal Einatc4bccd32018-09-12 00:49:13 +0300573/*[clinic input]
574@classmethod
575itertools.teedataobject.__new__
576 iterable as it: object
577 values: object(subclass_of='&PyList_Type')
578 next: object
579 /
580Data container common to multiple tee objects.
581[clinic start generated code]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000582
583static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300584itertools_teedataobject_impl(PyTypeObject *type, PyObject *it,
585 PyObject *values, PyObject *next)
586/*[clinic end generated code: output=3343ceb07e08df5e input=be60f2fabd2b72ba]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000587{
588 teedataobject *tdo;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000589 Py_ssize_t i, len;
590
591 assert(type == &teedataobject_type);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000592
593 tdo = (teedataobject *)teedataobject_newinternal(it);
594 if (!tdo)
595 return NULL;
596
597 len = PyList_GET_SIZE(values);
598 if (len > LINKCELLS)
599 goto err;
600 for (i=0; i<len; i++) {
601 tdo->values[i] = PyList_GET_ITEM(values, i);
602 Py_INCREF(tdo->values[i]);
603 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200604 /* len <= LINKCELLS < INT_MAX */
605 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000606
607 if (len == LINKCELLS) {
608 if (next != Py_None) {
609 if (Py_TYPE(next) != &teedataobject_type)
610 goto err;
611 assert(tdo->nextlink == NULL);
612 Py_INCREF(next);
613 tdo->nextlink = next;
614 }
615 } else {
616 if (next != Py_None)
617 goto err; /* shouldn't have a next if we are not full */
618 }
619 return (PyObject*)tdo;
620
621err:
622 Py_XDECREF(tdo);
623 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
624 return NULL;
625}
626
627static PyMethodDef teedataobject_methods[] = {
628 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
629 reduce_doc},
630 {NULL, NULL} /* sentinel */
631};
632
Raymond Hettingerad983e72003-11-12 14:32:26 +0000633static PyTypeObject teedataobject_type = {
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700634 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000635 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000636 sizeof(teedataobject), /* tp_basicsize */
637 0, /* tp_itemsize */
638 /* methods */
639 (destructor)teedataobject_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200640 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000641 0, /* tp_getattr */
642 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200643 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000644 0, /* tp_repr */
645 0, /* tp_as_number */
646 0, /* tp_as_sequence */
647 0, /* tp_as_mapping */
648 0, /* tp_hash */
649 0, /* tp_call */
650 0, /* tp_str */
651 PyObject_GenericGetAttr, /* tp_getattro */
652 0, /* tp_setattro */
653 0, /* tp_as_buffer */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700654 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300655 itertools_teedataobject__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000656 (traverseproc)teedataobject_traverse, /* tp_traverse */
657 (inquiry)teedataobject_clear, /* tp_clear */
658 0, /* tp_richcompare */
659 0, /* tp_weaklistoffset */
660 0, /* tp_iter */
661 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000662 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000663 0, /* tp_members */
664 0, /* tp_getset */
665 0, /* tp_base */
666 0, /* tp_dict */
667 0, /* tp_descr_get */
668 0, /* tp_descr_set */
669 0, /* tp_dictoffset */
670 0, /* tp_init */
671 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +0300672 itertools_teedataobject, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000673 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000674};
675
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000676
677static PyTypeObject tee_type;
678
679static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000680tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000681{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000682 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000683
HongWeipengfa220ec2019-08-29 11:39:25 +0800684 if (to->thread_id != PyThread_get_thread_ident()) {
685 PyErr_SetString(PyExc_RuntimeError,
686 "tee() iterator can not be consumed from different threads.");
687 return NULL;
688 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000689 if (to->index >= LINKCELLS) {
690 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200691 if (link == NULL)
692 return NULL;
Serhiy Storchaka57a01d32016-04-10 18:05:40 +0300693 Py_SETREF(to->dataobj, (teedataobject *)link);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000694 to->index = 0;
695 }
696 value = teedataobject_getitem(to->dataobj, to->index);
697 if (value == NULL)
698 return NULL;
699 to->index++;
700 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000701}
702
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000703static int
704tee_traverse(teeobject *to, visitproc visit, void *arg)
705{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000706 Py_VISIT((PyObject *)to->dataobj);
707 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000708}
709
Raymond Hettingerad983e72003-11-12 14:32:26 +0000710static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530711tee_copy(teeobject *to, PyObject *Py_UNUSED(ignored))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000712{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000713 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000714
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000715 newto = PyObject_GC_New(teeobject, &tee_type);
716 if (newto == NULL)
717 return NULL;
718 Py_INCREF(to->dataobj);
719 newto->dataobj = to->dataobj;
720 newto->index = to->index;
721 newto->weakreflist = NULL;
HongWeipengfa220ec2019-08-29 11:39:25 +0800722 newto->thread_id = to->thread_id;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000723 PyObject_GC_Track(newto);
724 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000725}
726
727PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
728
729static PyObject *
730tee_fromiterable(PyObject *iterable)
731{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000732 teeobject *to;
733 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000734
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000735 it = PyObject_GetIter(iterable);
736 if (it == NULL)
737 return NULL;
738 if (PyObject_TypeCheck(it, &tee_type)) {
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530739 to = (teeobject *)tee_copy((teeobject *)it, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000740 goto done;
741 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000742
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000743 to = PyObject_GC_New(teeobject, &tee_type);
744 if (to == NULL)
745 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000746 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000747 if (!to->dataobj) {
748 PyObject_GC_Del(to);
749 to = NULL;
750 goto done;
751 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000752
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000753 to->index = 0;
754 to->weakreflist = NULL;
HongWeipengfa220ec2019-08-29 11:39:25 +0800755 to->thread_id = PyThread_get_thread_ident();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000756 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000757done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000758 Py_XDECREF(it);
759 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000760}
761
Tal Einatc4bccd32018-09-12 00:49:13 +0300762/*[clinic input]
763@classmethod
764itertools._tee.__new__
765 iterable: object
766 /
767Iterator wrapped to make it copyable.
768[clinic start generated code]*/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000769
Tal Einatc4bccd32018-09-12 00:49:13 +0300770static PyObject *
771itertools__tee_impl(PyTypeObject *type, PyObject *iterable)
772/*[clinic end generated code: output=b02d3fd26c810c3f input=adc0779d2afe37a2]*/
773{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000774 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000775}
776
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000777static int
778tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000779{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000780 if (to->weakreflist != NULL)
781 PyObject_ClearWeakRefs((PyObject *) to);
782 Py_CLEAR(to->dataobj);
783 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000784}
785
786static void
787tee_dealloc(teeobject *to)
788{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000789 PyObject_GC_UnTrack(to);
790 tee_clear(to);
791 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000792}
793
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000794static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530795tee_reduce(teeobject *to, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000796{
797 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
798}
799
800static PyObject *
801tee_setstate(teeobject *to, PyObject *state)
802{
803 teedataobject *tdo;
804 int index;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300805 if (!PyTuple_Check(state)) {
806 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000807 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300808 }
809 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {
810 return NULL;
811 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000812 if (index < 0 || index > LINKCELLS) {
813 PyErr_SetString(PyExc_ValueError, "Index out of range");
814 return NULL;
815 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200816 Py_INCREF(tdo);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300817 Py_XSETREF(to->dataobj, tdo);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000818 to->index = index;
819 Py_RETURN_NONE;
820}
821
Raymond Hettingerad983e72003-11-12 14:32:26 +0000822static PyMethodDef tee_methods[] = {
Raymond Hettingera6a2d442015-08-16 14:38:07 -0700823 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
824 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
825 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000826 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000827};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000828
829static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000830 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000831 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000832 sizeof(teeobject), /* tp_basicsize */
833 0, /* tp_itemsize */
834 /* methods */
835 (destructor)tee_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200836 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000837 0, /* tp_getattr */
838 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200839 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000840 0, /* tp_repr */
841 0, /* tp_as_number */
842 0, /* tp_as_sequence */
843 0, /* tp_as_mapping */
844 0, /* tp_hash */
845 0, /* tp_call */
846 0, /* tp_str */
847 0, /* tp_getattro */
848 0, /* tp_setattro */
849 0, /* tp_as_buffer */
850 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300851 itertools__tee__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000852 (traverseproc)tee_traverse, /* tp_traverse */
853 (inquiry)tee_clear, /* tp_clear */
854 0, /* tp_richcompare */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700855 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000856 PyObject_SelfIter, /* tp_iter */
857 (iternextfunc)tee_next, /* tp_iternext */
858 tee_methods, /* tp_methods */
859 0, /* tp_members */
860 0, /* tp_getset */
861 0, /* tp_base */
862 0, /* tp_dict */
863 0, /* tp_descr_get */
864 0, /* tp_descr_set */
865 0, /* tp_dictoffset */
866 0, /* tp_init */
867 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +0300868 itertools__tee, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000869 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000870};
871
Tal Einatc4bccd32018-09-12 00:49:13 +0300872/*[clinic input]
873itertools.tee
874 iterable: object
875 n: Py_ssize_t = 2
876 /
877Returns a tuple of n independent iterators.
878[clinic start generated code]*/
879
Raymond Hettingerad983e72003-11-12 14:32:26 +0000880static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300881itertools_tee_impl(PyObject *module, PyObject *iterable, Py_ssize_t n)
882/*[clinic end generated code: output=1c64519cd859c2f0 input=c99a1472c425d66d]*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000883{
Tal Einatc4bccd32018-09-12 00:49:13 +0300884 Py_ssize_t i;
885 PyObject *it, *copyable, *copyfunc, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200886 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000887
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000888 if (n < 0) {
889 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
890 return NULL;
891 }
892 result = PyTuple_New(n);
893 if (result == NULL)
894 return NULL;
895 if (n == 0)
896 return result;
897 it = PyObject_GetIter(iterable);
898 if (it == NULL) {
899 Py_DECREF(result);
900 return NULL;
901 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200902
Serhiy Storchakaf320be72018-01-25 10:49:40 +0200903 if (_PyObject_LookupAttrId(it, &PyId___copy__, &copyfunc) < 0) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200904 Py_DECREF(it);
905 Py_DECREF(result);
906 return NULL;
907 }
Serhiy Storchakaf320be72018-01-25 10:49:40 +0200908 if (copyfunc != NULL) {
909 copyable = it;
910 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200911 else {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000912 copyable = tee_fromiterable(it);
913 Py_DECREF(it);
914 if (copyable == NULL) {
915 Py_DECREF(result);
916 return NULL;
917 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200918 copyfunc = _PyObject_GetAttrId(copyable, &PyId___copy__);
919 if (copyfunc == NULL) {
920 Py_DECREF(copyable);
921 Py_DECREF(result);
922 return NULL;
923 }
924 }
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200925
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200926 PyTuple_SET_ITEM(result, 0, copyable);
927 for (i = 1; i < n; i++) {
928 copyable = _PyObject_CallNoArg(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000929 if (copyable == NULL) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200930 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000931 Py_DECREF(result);
932 return NULL;
933 }
934 PyTuple_SET_ITEM(result, i, copyable);
935 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200936 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000937 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000938}
939
Raymond Hettingerad983e72003-11-12 14:32:26 +0000940
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700941/* cycle object **************************************************************/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000942
943typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000944 PyObject_HEAD
945 PyObject *it;
946 PyObject *saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700947 Py_ssize_t index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000948 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000949} cycleobject;
950
Raymond Hettinger1d7a3482003-07-14 07:07:12 +0000951static PyTypeObject cycle_type;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000952
Tal Einatc4bccd32018-09-12 00:49:13 +0300953/*[clinic input]
954@classmethod
955itertools.cycle.__new__
956 iterable: object
957 /
958Return elements from the iterable until it is exhausted. Then repeat the sequence indefinitely.
959[clinic start generated code]*/
960
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000961static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300962itertools_cycle_impl(PyTypeObject *type, PyObject *iterable)
963/*[clinic end generated code: output=f60e5ec17a45b35c input=9d1d84bcf66e908b]*/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000964{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000965 PyObject *it;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000966 PyObject *saved;
967 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000968
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000969 /* Get iterator. */
970 it = PyObject_GetIter(iterable);
971 if (it == NULL)
972 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000973
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000974 saved = PyList_New(0);
975 if (saved == NULL) {
976 Py_DECREF(it);
977 return NULL;
978 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000979
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000980 /* create cycleobject structure */
981 lz = (cycleobject *)type->tp_alloc(type, 0);
982 if (lz == NULL) {
983 Py_DECREF(it);
984 Py_DECREF(saved);
985 return NULL;
986 }
987 lz->it = it;
988 lz->saved = saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700989 lz->index = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000990 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000991
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000992 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000993}
994
995static void
996cycle_dealloc(cycleobject *lz)
997{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000998 PyObject_GC_UnTrack(lz);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000999 Py_XDECREF(lz->it);
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001000 Py_XDECREF(lz->saved);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001001 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001002}
1003
1004static int
1005cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
1006{
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001007 if (lz->it)
1008 Py_VISIT(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001009 Py_VISIT(lz->saved);
1010 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001011}
1012
1013static PyObject *
1014cycle_next(cycleobject *lz)
1015{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001016 PyObject *item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001017
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001018 if (lz->it != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001019 item = PyIter_Next(lz->it);
1020 if (item != NULL) {
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001021 if (lz->firstpass)
1022 return item;
1023 if (PyList_Append(lz->saved, item)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001024 Py_DECREF(item);
1025 return NULL;
1026 }
1027 return item;
1028 }
Raymond Hettinger98958fe2015-08-15 15:09:30 -07001029 /* Note: StopIteration is already cleared by PyIter_Next() */
1030 if (PyErr_Occurred())
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001031 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001032 Py_CLEAR(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001033 }
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001034 if (PyList_GET_SIZE(lz->saved) == 0)
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001035 return NULL;
1036 item = PyList_GET_ITEM(lz->saved, lz->index);
1037 lz->index++;
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001038 if (lz->index >= PyList_GET_SIZE(lz->saved))
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001039 lz->index = 0;
1040 Py_INCREF(item);
1041 return item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001042}
1043
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001044static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301045cycle_reduce(cycleobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001046{
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001047 /* Create a new cycle with the iterator tuple, then set the saved state */
1048 if (lz->it == NULL) {
1049 PyObject *it = PyObject_GetIter(lz->saved);
1050 if (it == NULL)
1051 return NULL;
1052 if (lz->index != 0) {
1053 _Py_IDENTIFIER(__setstate__);
1054 PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
1055 "n", lz->index);
1056 if (res == NULL) {
1057 Py_DECREF(it);
1058 return NULL;
1059 }
1060 Py_DECREF(res);
1061 }
Serhiy Storchaka1f21eaa2019-09-01 12:16:51 +03001062 return Py_BuildValue("O(N)(OO)", Py_TYPE(lz), it, lz->saved, Py_True);
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001063 }
Serhiy Storchaka1f21eaa2019-09-01 12:16:51 +03001064 return Py_BuildValue("O(O)(OO)", Py_TYPE(lz), lz->it, lz->saved,
1065 lz->firstpass ? Py_True : Py_False);
Raymond Hettingerc39786d2015-08-15 15:16:12 -07001066}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001067
1068static PyObject *
1069cycle_setstate(cycleobject *lz, PyObject *state)
1070{
1071 PyObject *saved=NULL;
1072 int firstpass;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001073 if (!PyTuple_Check(state)) {
1074 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001075 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001076 }
1077 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
1078 return NULL;
1079 }
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001080 Py_INCREF(saved);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001081 Py_XSETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001082 lz->firstpass = firstpass != 0;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001083 lz->index = 0;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001084 Py_RETURN_NONE;
1085}
1086
1087static PyMethodDef cycle_methods[] = {
1088 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
1089 reduce_doc},
1090 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
1091 setstate_doc},
1092 {NULL, NULL} /* sentinel */
1093};
1094
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001095static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001096 PyVarObject_HEAD_INIT(NULL, 0)
1097 "itertools.cycle", /* tp_name */
1098 sizeof(cycleobject), /* tp_basicsize */
1099 0, /* tp_itemsize */
1100 /* methods */
1101 (destructor)cycle_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001102 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001103 0, /* tp_getattr */
1104 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001105 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001106 0, /* tp_repr */
1107 0, /* tp_as_number */
1108 0, /* tp_as_sequence */
1109 0, /* tp_as_mapping */
1110 0, /* tp_hash */
1111 0, /* tp_call */
1112 0, /* tp_str */
1113 PyObject_GenericGetAttr, /* tp_getattro */
1114 0, /* tp_setattro */
1115 0, /* tp_as_buffer */
1116 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1117 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001118 itertools_cycle__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001119 (traverseproc)cycle_traverse, /* tp_traverse */
1120 0, /* tp_clear */
1121 0, /* tp_richcompare */
1122 0, /* tp_weaklistoffset */
1123 PyObject_SelfIter, /* tp_iter */
1124 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001125 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001126 0, /* tp_members */
1127 0, /* tp_getset */
1128 0, /* tp_base */
1129 0, /* tp_dict */
1130 0, /* tp_descr_get */
1131 0, /* tp_descr_set */
1132 0, /* tp_dictoffset */
1133 0, /* tp_init */
1134 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001135 itertools_cycle, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001136 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001137};
1138
1139
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001140/* dropwhile object **********************************************************/
1141
1142typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001143 PyObject_HEAD
1144 PyObject *func;
1145 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001146 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001147} dropwhileobject;
1148
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001149static PyTypeObject dropwhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001150
Tal Einatc4bccd32018-09-12 00:49:13 +03001151/*[clinic input]
1152@classmethod
1153itertools.dropwhile.__new__
1154 predicate as func: object
1155 iterable as seq: object
1156 /
1157Drop items from the iterable while predicate(item) is true.
1158
1159Afterwards, return every element until the iterable is exhausted.
1160[clinic start generated code]*/
1161
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001162static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001163itertools_dropwhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1164/*[clinic end generated code: output=92f9d0d89af149e4 input=d39737147c9f0a26]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001165{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001166 PyObject *it;
1167 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001168
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001169 /* Get iterator. */
1170 it = PyObject_GetIter(seq);
1171 if (it == NULL)
1172 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001173
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001174 /* create dropwhileobject structure */
1175 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1176 if (lz == NULL) {
1177 Py_DECREF(it);
1178 return NULL;
1179 }
1180 Py_INCREF(func);
1181 lz->func = func;
1182 lz->it = it;
1183 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001184
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001185 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001186}
1187
1188static void
1189dropwhile_dealloc(dropwhileobject *lz)
1190{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001191 PyObject_GC_UnTrack(lz);
1192 Py_XDECREF(lz->func);
1193 Py_XDECREF(lz->it);
1194 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001195}
1196
1197static int
1198dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1199{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001200 Py_VISIT(lz->it);
1201 Py_VISIT(lz->func);
1202 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001203}
1204
1205static PyObject *
1206dropwhile_next(dropwhileobject *lz)
1207{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001208 PyObject *item, *good;
1209 PyObject *it = lz->it;
1210 long ok;
1211 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001212
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001213 iternext = *Py_TYPE(it)->tp_iternext;
1214 for (;;) {
1215 item = iternext(it);
1216 if (item == NULL)
1217 return NULL;
1218 if (lz->start == 1)
1219 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001220
Jeroen Demeyer196a5302019-07-04 12:31:34 +02001221 good = _PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001222 if (good == NULL) {
1223 Py_DECREF(item);
1224 return NULL;
1225 }
1226 ok = PyObject_IsTrue(good);
1227 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001228 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001229 lz->start = 1;
1230 return item;
1231 }
1232 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001233 if (ok < 0)
1234 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001235 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001236}
1237
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001238static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301239dropwhile_reduce(dropwhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001240{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001241 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 +00001242}
1243
1244static PyObject *
1245dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1246{
1247 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001248 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001249 return NULL;
1250 lz->start = start;
1251 Py_RETURN_NONE;
1252}
1253
1254static PyMethodDef dropwhile_methods[] = {
1255 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1256 reduce_doc},
1257 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1258 setstate_doc},
1259 {NULL, NULL} /* sentinel */
1260};
1261
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001262static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001263 PyVarObject_HEAD_INIT(NULL, 0)
1264 "itertools.dropwhile", /* tp_name */
1265 sizeof(dropwhileobject), /* tp_basicsize */
1266 0, /* tp_itemsize */
1267 /* methods */
1268 (destructor)dropwhile_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001269 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001270 0, /* tp_getattr */
1271 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001272 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001273 0, /* tp_repr */
1274 0, /* tp_as_number */
1275 0, /* tp_as_sequence */
1276 0, /* tp_as_mapping */
1277 0, /* tp_hash */
1278 0, /* tp_call */
1279 0, /* tp_str */
1280 PyObject_GenericGetAttr, /* tp_getattro */
1281 0, /* tp_setattro */
1282 0, /* tp_as_buffer */
1283 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1284 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001285 itertools_dropwhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001286 (traverseproc)dropwhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001287 0, /* tp_clear */
1288 0, /* tp_richcompare */
1289 0, /* tp_weaklistoffset */
1290 PyObject_SelfIter, /* tp_iter */
1291 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001292 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001293 0, /* tp_members */
1294 0, /* tp_getset */
1295 0, /* tp_base */
1296 0, /* tp_dict */
1297 0, /* tp_descr_get */
1298 0, /* tp_descr_set */
1299 0, /* tp_dictoffset */
1300 0, /* tp_init */
1301 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001302 itertools_dropwhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001303 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001304};
1305
1306
1307/* takewhile object **********************************************************/
1308
1309typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001310 PyObject_HEAD
1311 PyObject *func;
1312 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001313 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001314} takewhileobject;
1315
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001316static PyTypeObject takewhile_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001317
Tal Einatc4bccd32018-09-12 00:49:13 +03001318/*[clinic input]
1319@classmethod
1320itertools.takewhile.__new__
1321 predicate as func: object
1322 iterable as seq: object
1323 /
1324Return successive entries from an iterable as long as the predicate evaluates to true for each entry.
1325[clinic start generated code]*/
1326
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001327static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001328itertools_takewhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1329/*[clinic end generated code: output=bb179ea7864e2ef6 input=ba5255f7519aa119]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001330{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001331 PyObject *it;
1332 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001333
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001334 /* Get iterator. */
1335 it = PyObject_GetIter(seq);
1336 if (it == NULL)
1337 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001338
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001339 /* create takewhileobject structure */
1340 lz = (takewhileobject *)type->tp_alloc(type, 0);
1341 if (lz == NULL) {
1342 Py_DECREF(it);
1343 return NULL;
1344 }
1345 Py_INCREF(func);
1346 lz->func = func;
1347 lz->it = it;
1348 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001349
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001350 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001351}
1352
1353static void
1354takewhile_dealloc(takewhileobject *lz)
1355{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001356 PyObject_GC_UnTrack(lz);
1357 Py_XDECREF(lz->func);
1358 Py_XDECREF(lz->it);
1359 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001360}
1361
1362static int
1363takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1364{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001365 Py_VISIT(lz->it);
1366 Py_VISIT(lz->func);
1367 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001368}
1369
1370static PyObject *
1371takewhile_next(takewhileobject *lz)
1372{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001373 PyObject *item, *good;
1374 PyObject *it = lz->it;
1375 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001376
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001377 if (lz->stop == 1)
1378 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001379
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001380 item = (*Py_TYPE(it)->tp_iternext)(it);
1381 if (item == NULL)
1382 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001383
Jeroen Demeyer196a5302019-07-04 12:31:34 +02001384 good = _PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001385 if (good == NULL) {
1386 Py_DECREF(item);
1387 return NULL;
1388 }
1389 ok = PyObject_IsTrue(good);
1390 Py_DECREF(good);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001391 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001392 return item;
1393 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001394 if (ok == 0)
1395 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001396 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001397}
1398
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001399static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301400takewhile_reduce(takewhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001401{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001402 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 +00001403}
1404
1405static PyObject *
1406takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1407{
1408 int stop = PyObject_IsTrue(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001409
Antoine Pitrou721738f2012-08-15 23:20:39 +02001410 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001411 return NULL;
1412 lz->stop = stop;
1413 Py_RETURN_NONE;
1414}
1415
1416static PyMethodDef takewhile_reduce_methods[] = {
1417 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1418 reduce_doc},
1419 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1420 setstate_doc},
1421 {NULL, NULL} /* sentinel */
1422};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001423
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001424static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001425 PyVarObject_HEAD_INIT(NULL, 0)
1426 "itertools.takewhile", /* tp_name */
1427 sizeof(takewhileobject), /* tp_basicsize */
1428 0, /* tp_itemsize */
1429 /* methods */
1430 (destructor)takewhile_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001431 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001432 0, /* tp_getattr */
1433 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001434 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001435 0, /* tp_repr */
1436 0, /* tp_as_number */
1437 0, /* tp_as_sequence */
1438 0, /* tp_as_mapping */
1439 0, /* tp_hash */
1440 0, /* tp_call */
1441 0, /* tp_str */
1442 PyObject_GenericGetAttr, /* tp_getattro */
1443 0, /* tp_setattro */
1444 0, /* tp_as_buffer */
1445 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1446 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001447 itertools_takewhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001448 (traverseproc)takewhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001449 0, /* tp_clear */
1450 0, /* tp_richcompare */
1451 0, /* tp_weaklistoffset */
1452 PyObject_SelfIter, /* tp_iter */
1453 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001454 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001455 0, /* tp_members */
1456 0, /* tp_getset */
1457 0, /* tp_base */
1458 0, /* tp_dict */
1459 0, /* tp_descr_get */
1460 0, /* tp_descr_set */
1461 0, /* tp_dictoffset */
1462 0, /* tp_init */
1463 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001464 itertools_takewhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001465 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001466};
1467
1468
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001469/* islice object *************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001470
1471typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001472 PyObject_HEAD
1473 PyObject *it;
1474 Py_ssize_t next;
1475 Py_ssize_t stop;
1476 Py_ssize_t step;
1477 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001478} isliceobject;
1479
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001480static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001481
1482static PyObject *
1483islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1484{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001485 PyObject *seq;
1486 Py_ssize_t start=0, stop=-1, step=1;
1487 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1488 Py_ssize_t numargs;
1489 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001490
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001491 if (type == &islice_type && !_PyArg_NoKeywords("islice", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001492 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001493
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001494 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1495 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001496
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001497 numargs = PyTuple_Size(args);
1498 if (numargs == 2) {
1499 if (a1 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001500 stop = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001501 if (stop == -1) {
1502 if (PyErr_Occurred())
1503 PyErr_Clear();
1504 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001505 "Stop argument for islice() must be None or "
1506 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001507 return NULL;
1508 }
1509 }
1510 } else {
1511 if (a1 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001512 start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001513 if (start == -1 && PyErr_Occurred())
1514 PyErr_Clear();
1515 if (a2 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001516 stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001517 if (stop == -1) {
1518 if (PyErr_Occurred())
1519 PyErr_Clear();
1520 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001521 "Stop argument for islice() must be None or "
1522 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001523 return NULL;
1524 }
1525 }
1526 }
1527 if (start<0 || stop<-1) {
1528 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001529 "Indices for islice() must be None or "
1530 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001531 return NULL;
1532 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001533
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001534 if (a3 != NULL) {
1535 if (a3 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001536 step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001537 if (step == -1 && PyErr_Occurred())
1538 PyErr_Clear();
1539 }
1540 if (step<1) {
1541 PyErr_SetString(PyExc_ValueError,
1542 "Step for islice() must be a positive integer or None.");
1543 return NULL;
1544 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001545
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001546 /* Get iterator. */
1547 it = PyObject_GetIter(seq);
1548 if (it == NULL)
1549 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001550
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001551 /* create isliceobject structure */
1552 lz = (isliceobject *)type->tp_alloc(type, 0);
1553 if (lz == NULL) {
1554 Py_DECREF(it);
1555 return NULL;
1556 }
1557 lz->it = it;
1558 lz->next = start;
1559 lz->stop = stop;
1560 lz->step = step;
1561 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001562
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001563 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001564}
1565
1566static void
1567islice_dealloc(isliceobject *lz)
1568{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001569 PyObject_GC_UnTrack(lz);
1570 Py_XDECREF(lz->it);
1571 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001572}
1573
1574static int
1575islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1576{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001577 Py_VISIT(lz->it);
1578 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001579}
1580
1581static PyObject *
1582islice_next(isliceobject *lz)
1583{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001584 PyObject *item;
1585 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001586 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001587 Py_ssize_t oldnext;
1588 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001589
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001590 if (it == NULL)
1591 return NULL;
1592
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001593 iternext = *Py_TYPE(it)->tp_iternext;
1594 while (lz->cnt < lz->next) {
1595 item = iternext(it);
1596 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001597 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001598 Py_DECREF(item);
1599 lz->cnt++;
1600 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001601 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001602 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001603 item = iternext(it);
1604 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001605 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001606 lz->cnt++;
1607 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001608 /* The (size_t) cast below avoids the danger of undefined
1609 behaviour from signed integer overflow. */
1610 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001611 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1612 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001613 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001614
1615empty:
1616 Py_CLEAR(lz->it);
1617 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001618}
1619
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001620static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301621islice_reduce(isliceobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001622{
1623 /* When unpickled, generate a new object with the same bounds,
1624 * then 'setstate' with the next and count
1625 */
1626 PyObject *stop;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001627
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001628 if (lz->it == NULL) {
1629 PyObject *empty_list;
1630 PyObject *empty_it;
1631 empty_list = PyList_New(0);
1632 if (empty_list == NULL)
1633 return NULL;
1634 empty_it = PyObject_GetIter(empty_list);
1635 Py_DECREF(empty_list);
1636 if (empty_it == NULL)
1637 return NULL;
1638 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1639 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001640 if (lz->stop == -1) {
1641 stop = Py_None;
1642 Py_INCREF(stop);
1643 } else {
1644 stop = PyLong_FromSsize_t(lz->stop);
1645 if (stop == NULL)
1646 return NULL;
1647 }
1648 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1649 lz->it, lz->next, stop, lz->step,
1650 lz->cnt);
1651}
1652
1653static PyObject *
1654islice_setstate(isliceobject *lz, PyObject *state)
1655{
1656 Py_ssize_t cnt = PyLong_AsSsize_t(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001657
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001658 if (cnt == -1 && PyErr_Occurred())
1659 return NULL;
1660 lz->cnt = cnt;
1661 Py_RETURN_NONE;
1662}
1663
1664static PyMethodDef islice_methods[] = {
1665 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1666 reduce_doc},
1667 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1668 setstate_doc},
1669 {NULL, NULL} /* sentinel */
1670};
1671
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001672PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001673"islice(iterable, stop) --> islice object\n\
1674islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001675\n\
1676Return an iterator whose next() method returns selected values from an\n\
1677iterable. If start is specified, will skip all preceding elements;\n\
1678otherwise, start defaults to zero. Step defaults to one. If\n\
oldkaa0735f2018-02-02 16:52:55 +08001679specified as another value, step determines how many values are\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001680skipped between successive calls. Works like a slice() on a list\n\
1681but returns an iterator.");
1682
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001683static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001684 PyVarObject_HEAD_INIT(NULL, 0)
1685 "itertools.islice", /* tp_name */
1686 sizeof(isliceobject), /* tp_basicsize */
1687 0, /* tp_itemsize */
1688 /* methods */
1689 (destructor)islice_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001690 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001691 0, /* tp_getattr */
1692 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001693 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001694 0, /* tp_repr */
1695 0, /* tp_as_number */
1696 0, /* tp_as_sequence */
1697 0, /* tp_as_mapping */
1698 0, /* tp_hash */
1699 0, /* tp_call */
1700 0, /* tp_str */
1701 PyObject_GenericGetAttr, /* tp_getattro */
1702 0, /* tp_setattro */
1703 0, /* tp_as_buffer */
1704 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1705 Py_TPFLAGS_BASETYPE, /* tp_flags */
1706 islice_doc, /* tp_doc */
1707 (traverseproc)islice_traverse, /* tp_traverse */
1708 0, /* tp_clear */
1709 0, /* tp_richcompare */
1710 0, /* tp_weaklistoffset */
1711 PyObject_SelfIter, /* tp_iter */
1712 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001713 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001714 0, /* tp_members */
1715 0, /* tp_getset */
1716 0, /* tp_base */
1717 0, /* tp_dict */
1718 0, /* tp_descr_get */
1719 0, /* tp_descr_set */
1720 0, /* tp_dictoffset */
1721 0, /* tp_init */
1722 0, /* tp_alloc */
1723 islice_new, /* tp_new */
1724 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001725};
1726
1727
1728/* starmap object ************************************************************/
1729
1730typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001731 PyObject_HEAD
1732 PyObject *func;
1733 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001734} starmapobject;
1735
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001736static PyTypeObject starmap_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001737
Tal Einatc4bccd32018-09-12 00:49:13 +03001738/*[clinic input]
1739@classmethod
1740itertools.starmap.__new__
1741 function as func: object
1742 iterable as seq: object
1743 /
1744Return an iterator whose values are returned from the function evaluated with an argument tuple taken from the given sequence.
1745[clinic start generated code]*/
1746
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001747static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001748itertools_starmap_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1749/*[clinic end generated code: output=79eeb81d452c6e8d input=844766df6a0d4dad]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001750{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001751 PyObject *it;
1752 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001753
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001754 /* Get iterator. */
1755 it = PyObject_GetIter(seq);
1756 if (it == NULL)
1757 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001758
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001759 /* create starmapobject structure */
1760 lz = (starmapobject *)type->tp_alloc(type, 0);
1761 if (lz == NULL) {
1762 Py_DECREF(it);
1763 return NULL;
1764 }
1765 Py_INCREF(func);
1766 lz->func = func;
1767 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001768
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001769 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001770}
1771
1772static void
1773starmap_dealloc(starmapobject *lz)
1774{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001775 PyObject_GC_UnTrack(lz);
1776 Py_XDECREF(lz->func);
1777 Py_XDECREF(lz->it);
1778 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001779}
1780
1781static int
1782starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1783{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001784 Py_VISIT(lz->it);
1785 Py_VISIT(lz->func);
1786 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001787}
1788
1789static PyObject *
1790starmap_next(starmapobject *lz)
1791{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001792 PyObject *args;
1793 PyObject *result;
1794 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001795
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001796 args = (*Py_TYPE(it)->tp_iternext)(it);
1797 if (args == NULL)
1798 return NULL;
1799 if (!PyTuple_CheckExact(args)) {
1800 PyObject *newargs = PySequence_Tuple(args);
1801 Py_DECREF(args);
1802 if (newargs == NULL)
1803 return NULL;
1804 args = newargs;
1805 }
1806 result = PyObject_Call(lz->func, args, NULL);
1807 Py_DECREF(args);
1808 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001809}
1810
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001811static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301812starmap_reduce(starmapobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001813{
1814 /* Just pickle the iterator */
1815 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1816}
1817
1818static PyMethodDef starmap_methods[] = {
1819 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1820 reduce_doc},
1821 {NULL, NULL} /* sentinel */
1822};
1823
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001824static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001825 PyVarObject_HEAD_INIT(NULL, 0)
1826 "itertools.starmap", /* tp_name */
1827 sizeof(starmapobject), /* tp_basicsize */
1828 0, /* tp_itemsize */
1829 /* methods */
1830 (destructor)starmap_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001831 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001832 0, /* tp_getattr */
1833 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001834 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001835 0, /* tp_repr */
1836 0, /* tp_as_number */
1837 0, /* tp_as_sequence */
1838 0, /* tp_as_mapping */
1839 0, /* tp_hash */
1840 0, /* tp_call */
1841 0, /* tp_str */
1842 PyObject_GenericGetAttr, /* tp_getattro */
1843 0, /* tp_setattro */
1844 0, /* tp_as_buffer */
1845 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1846 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001847 itertools_starmap__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001848 (traverseproc)starmap_traverse, /* tp_traverse */
1849 0, /* tp_clear */
1850 0, /* tp_richcompare */
1851 0, /* tp_weaklistoffset */
1852 PyObject_SelfIter, /* tp_iter */
1853 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001854 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001855 0, /* tp_members */
1856 0, /* tp_getset */
1857 0, /* tp_base */
1858 0, /* tp_dict */
1859 0, /* tp_descr_get */
1860 0, /* tp_descr_set */
1861 0, /* tp_dictoffset */
1862 0, /* tp_init */
1863 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001864 itertools_starmap, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001865 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001866};
1867
1868
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001869/* chain object **************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001870
1871typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001872 PyObject_HEAD
1873 PyObject *source; /* Iterator over input iterables */
1874 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001875} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001876
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001877static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001878
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001879static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001880chain_new_internal(PyTypeObject *type, PyObject *source)
1881{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001882 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001883
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001884 lz = (chainobject *)type->tp_alloc(type, 0);
1885 if (lz == NULL) {
1886 Py_DECREF(source);
1887 return NULL;
1888 }
1889
1890 lz->source = source;
1891 lz->active = NULL;
1892 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001893}
1894
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001895static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001896chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001897{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001898 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001899
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001900 if (type == &chain_type && !_PyArg_NoKeywords("chain", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001901 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001902
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001903 source = PyObject_GetIter(args);
1904 if (source == NULL)
1905 return NULL;
1906
1907 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001908}
1909
Tal Einatc4bccd32018-09-12 00:49:13 +03001910/*[clinic input]
1911@classmethod
1912itertools.chain.from_iterable
1913 iterable as arg: object
1914 /
1915Alternative chain() constructor taking a single iterable argument that evaluates lazily.
1916[clinic start generated code]*/
1917
Christian Heimesf16baeb2008-02-29 14:57:44 +00001918static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001919itertools_chain_from_iterable(PyTypeObject *type, PyObject *arg)
1920/*[clinic end generated code: output=667ae7a7f7b68654 input=72c39e3a2ca3be85]*/
Christian Heimesf16baeb2008-02-29 14:57:44 +00001921{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001922 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001923
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001924 source = PyObject_GetIter(arg);
1925 if (source == NULL)
1926 return NULL;
1927
1928 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001929}
1930
1931static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001932chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001933{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001934 PyObject_GC_UnTrack(lz);
1935 Py_XDECREF(lz->active);
1936 Py_XDECREF(lz->source);
1937 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001938}
1939
Raymond Hettinger2012f172003-02-07 05:32:58 +00001940static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001941chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001942{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001943 Py_VISIT(lz->source);
1944 Py_VISIT(lz->active);
1945 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001946}
1947
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001948static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001949chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001950{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001951 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001952
T. Wouters5466d4a2017-03-30 09:58:35 -07001953 /* lz->source is the iterator of iterables. If it's NULL, we've already
1954 * consumed them all. lz->active is the current iterator. If it's NULL,
1955 * we should grab a new one from lz->source. */
1956 while (lz->source != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001957 if (lz->active == NULL) {
T. Wouters5466d4a2017-03-30 09:58:35 -07001958 PyObject *iterable = PyIter_Next(lz->source);
1959 if (iterable == NULL) {
1960 Py_CLEAR(lz->source);
1961 return NULL; /* no more input sources */
1962 }
1963 lz->active = PyObject_GetIter(iterable);
1964 Py_DECREF(iterable);
1965 if (lz->active == NULL) {
1966 Py_CLEAR(lz->source);
1967 return NULL; /* input not iterable */
1968 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001969 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001970 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
1971 if (item != NULL)
1972 return item;
1973 if (PyErr_Occurred()) {
1974 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1975 PyErr_Clear();
1976 else
1977 return NULL; /* input raised an exception */
1978 }
1979 /* lz->active is consumed, try with the next iterable. */
1980 Py_CLEAR(lz->active);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001981 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001982 /* Everything had been consumed already. */
1983 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001984}
1985
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001986static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301987chain_reduce(chainobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001988{
1989 if (lz->source) {
1990 /* we can't pickle function objects (itertools.from_iterable) so
1991 * we must use setstate to replace the iterable. One day we
1992 * will fix pickling of functions
1993 */
1994 if (lz->active) {
1995 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1996 } else {
1997 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1998 }
1999 } else {
2000 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
2001 }
2002 return NULL;
2003}
2004
2005static PyObject *
2006chain_setstate(chainobject *lz, PyObject *state)
2007{
2008 PyObject *source, *active=NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002009
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002010 if (!PyTuple_Check(state)) {
2011 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002012 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002013 }
2014 if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
2015 return NULL;
2016 }
2017 if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
2018 PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
2019 return NULL;
2020 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002021
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002022 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002023 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002024 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002025 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002026 Py_RETURN_NONE;
2027}
2028
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002029PyDoc_STRVAR(chain_doc,
2030"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002031\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002032Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002033first iterable until it is exhausted, then elements from the next\n\
2034iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002035
Christian Heimesf16baeb2008-02-29 14:57:44 +00002036static PyMethodDef chain_methods[] = {
Tal Einatc4bccd32018-09-12 00:49:13 +03002037 ITERTOOLS_CHAIN_FROM_ITERABLE_METHODDEF
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002038 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
2039 reduce_doc},
2040 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
2041 setstate_doc},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07002042 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00002043};
2044
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002045static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002046 PyVarObject_HEAD_INIT(NULL, 0)
2047 "itertools.chain", /* tp_name */
2048 sizeof(chainobject), /* tp_basicsize */
2049 0, /* tp_itemsize */
2050 /* methods */
2051 (destructor)chain_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002052 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002053 0, /* tp_getattr */
2054 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002055 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002056 0, /* tp_repr */
2057 0, /* tp_as_number */
2058 0, /* tp_as_sequence */
2059 0, /* tp_as_mapping */
2060 0, /* tp_hash */
2061 0, /* tp_call */
2062 0, /* tp_str */
2063 PyObject_GenericGetAttr, /* tp_getattro */
2064 0, /* tp_setattro */
2065 0, /* tp_as_buffer */
2066 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2067 Py_TPFLAGS_BASETYPE, /* tp_flags */
2068 chain_doc, /* tp_doc */
2069 (traverseproc)chain_traverse, /* tp_traverse */
2070 0, /* tp_clear */
2071 0, /* tp_richcompare */
2072 0, /* tp_weaklistoffset */
2073 PyObject_SelfIter, /* tp_iter */
2074 (iternextfunc)chain_next, /* tp_iternext */
2075 chain_methods, /* tp_methods */
2076 0, /* tp_members */
2077 0, /* tp_getset */
2078 0, /* tp_base */
2079 0, /* tp_dict */
2080 0, /* tp_descr_get */
2081 0, /* tp_descr_set */
2082 0, /* tp_dictoffset */
2083 0, /* tp_init */
2084 0, /* tp_alloc */
2085 chain_new, /* tp_new */
2086 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002087};
2088
2089
Christian Heimesc3f30c42008-02-22 16:37:40 +00002090/* product object ************************************************************/
2091
2092typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002093 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002094 PyObject *pools; /* tuple of pool tuples */
2095 Py_ssize_t *indices; /* one index per pool */
2096 PyObject *result; /* most recently returned result tuple */
2097 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002098} productobject;
2099
2100static PyTypeObject product_type;
2101
2102static PyObject *
2103product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2104{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002105 productobject *lz;
2106 Py_ssize_t nargs, npools, repeat=1;
2107 PyObject *pools = NULL;
2108 Py_ssize_t *indices = NULL;
2109 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002110
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002111 if (kwds != NULL) {
2112 char *kwlist[] = {"repeat", 0};
2113 PyObject *tmpargs = PyTuple_New(0);
2114 if (tmpargs == NULL)
2115 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002116 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2117 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002118 Py_DECREF(tmpargs);
2119 return NULL;
2120 }
2121 Py_DECREF(tmpargs);
2122 if (repeat < 0) {
2123 PyErr_SetString(PyExc_ValueError,
2124 "repeat argument cannot be negative");
2125 return NULL;
2126 }
2127 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002128
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002129 assert(PyTuple_CheckExact(args));
2130 if (repeat == 0) {
2131 nargs = 0;
2132 } else {
2133 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002134 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002135 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2136 return NULL;
2137 }
2138 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002139 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002140
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002141 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002142 if (indices == NULL) {
2143 PyErr_NoMemory();
2144 goto error;
2145 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002146
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002147 pools = PyTuple_New(npools);
2148 if (pools == NULL)
2149 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002150
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002151 for (i=0; i < nargs ; ++i) {
2152 PyObject *item = PyTuple_GET_ITEM(args, i);
2153 PyObject *pool = PySequence_Tuple(item);
2154 if (pool == NULL)
2155 goto error;
2156 PyTuple_SET_ITEM(pools, i, pool);
2157 indices[i] = 0;
2158 }
2159 for ( ; i < npools; ++i) {
2160 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2161 Py_INCREF(pool);
2162 PyTuple_SET_ITEM(pools, i, pool);
2163 indices[i] = 0;
2164 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002165
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002166 /* create productobject structure */
2167 lz = (productobject *)type->tp_alloc(type, 0);
2168 if (lz == NULL)
2169 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002170
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002171 lz->pools = pools;
2172 lz->indices = indices;
2173 lz->result = NULL;
2174 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002175
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002176 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002177
2178error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002179 if (indices != NULL)
2180 PyMem_Free(indices);
2181 Py_XDECREF(pools);
2182 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002183}
2184
2185static void
2186product_dealloc(productobject *lz)
2187{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002188 PyObject_GC_UnTrack(lz);
2189 Py_XDECREF(lz->pools);
2190 Py_XDECREF(lz->result);
2191 if (lz->indices != NULL)
2192 PyMem_Free(lz->indices);
2193 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002194}
2195
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002196static PyObject *
2197product_sizeof(productobject *lz, void *unused)
2198{
2199 Py_ssize_t res;
2200
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002201 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002202 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2203 return PyLong_FromSsize_t(res);
2204}
2205
2206PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2207
Christian Heimesc3f30c42008-02-22 16:37:40 +00002208static int
2209product_traverse(productobject *lz, visitproc visit, void *arg)
2210{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002211 Py_VISIT(lz->pools);
2212 Py_VISIT(lz->result);
2213 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002214}
2215
2216static PyObject *
2217product_next(productobject *lz)
2218{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002219 PyObject *pool;
2220 PyObject *elem;
2221 PyObject *oldelem;
2222 PyObject *pools = lz->pools;
2223 PyObject *result = lz->result;
2224 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2225 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002226
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002227 if (lz->stopped)
2228 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002229
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002230 if (result == NULL) {
2231 /* On the first pass, return an initial tuple filled with the
2232 first element from each pool. */
2233 result = PyTuple_New(npools);
2234 if (result == NULL)
2235 goto empty;
2236 lz->result = result;
2237 for (i=0; i < npools; i++) {
2238 pool = PyTuple_GET_ITEM(pools, i);
2239 if (PyTuple_GET_SIZE(pool) == 0)
2240 goto empty;
2241 elem = PyTuple_GET_ITEM(pool, 0);
2242 Py_INCREF(elem);
2243 PyTuple_SET_ITEM(result, i, elem);
2244 }
2245 } else {
2246 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002247
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002248 /* Copy the previous result tuple or re-use it if available */
2249 if (Py_REFCNT(result) > 1) {
2250 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002251 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002252 if (result == NULL)
2253 goto empty;
2254 lz->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002255 Py_DECREF(old_result);
2256 }
2257 /* Now, we've got the only copy so we can update it in-place */
2258 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002259
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002260 /* Update the pool indices right-to-left. Only advance to the
2261 next pool when the previous one rolls-over */
2262 for (i=npools-1 ; i >= 0 ; i--) {
2263 pool = PyTuple_GET_ITEM(pools, i);
2264 indices[i]++;
2265 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2266 /* Roll-over and advance to next pool */
2267 indices[i] = 0;
2268 elem = PyTuple_GET_ITEM(pool, 0);
2269 Py_INCREF(elem);
2270 oldelem = PyTuple_GET_ITEM(result, i);
2271 PyTuple_SET_ITEM(result, i, elem);
2272 Py_DECREF(oldelem);
2273 } else {
2274 /* No rollover. Just increment and stop here. */
2275 elem = PyTuple_GET_ITEM(pool, indices[i]);
2276 Py_INCREF(elem);
2277 oldelem = PyTuple_GET_ITEM(result, i);
2278 PyTuple_SET_ITEM(result, i, elem);
2279 Py_DECREF(oldelem);
2280 break;
2281 }
2282 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002283
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002284 /* If i is negative, then the indices have all rolled-over
2285 and we're done. */
2286 if (i < 0)
2287 goto empty;
2288 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002289
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002290 Py_INCREF(result);
2291 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002292
2293empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002294 lz->stopped = 1;
2295 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002296}
2297
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002298static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302299product_reduce(productobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002300{
2301 if (lz->stopped) {
2302 return Py_BuildValue("O(())", Py_TYPE(lz));
2303 } else if (lz->result == NULL) {
2304 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2305 } else {
2306 PyObject *indices;
2307 Py_ssize_t n, i;
2308
2309 /* we must pickle the indices use them for setstate, and
2310 * additionally indicate that the iterator has started
2311 */
2312 n = PyTuple_GET_SIZE(lz->pools);
2313 indices = PyTuple_New(n);
2314 if (indices == NULL)
2315 return NULL;
2316 for (i=0; i<n; i++){
2317 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2318 if (!index) {
2319 Py_DECREF(indices);
2320 return NULL;
2321 }
2322 PyTuple_SET_ITEM(indices, i, index);
2323 }
2324 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2325 }
2326}
2327
2328static PyObject *
2329product_setstate(productobject *lz, PyObject *state)
2330{
2331 PyObject *result;
2332 Py_ssize_t n, i;
2333
2334 n = PyTuple_GET_SIZE(lz->pools);
2335 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2336 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2337 return NULL;
2338 }
2339 for (i=0; i<n; i++)
2340 {
2341 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2342 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002343 PyObject* pool;
2344 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002345 if (index < 0 && PyErr_Occurred())
2346 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002347 pool = PyTuple_GET_ITEM(lz->pools, i);
2348 poolsize = PyTuple_GET_SIZE(pool);
2349 if (poolsize == 0) {
2350 lz->stopped = 1;
2351 Py_RETURN_NONE;
2352 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002353 /* clamp the index */
2354 if (index < 0)
2355 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002356 else if (index > poolsize-1)
2357 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002358 lz->indices[i] = index;
2359 }
2360
2361 result = PyTuple_New(n);
2362 if (!result)
2363 return NULL;
2364 for (i=0; i<n; i++) {
2365 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2366 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2367 Py_INCREF(element);
2368 PyTuple_SET_ITEM(result, i, element);
2369 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002370 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002371 Py_RETURN_NONE;
2372}
2373
2374static PyMethodDef product_methods[] = {
2375 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2376 reduce_doc},
2377 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2378 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002379 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2380 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002381 {NULL, NULL} /* sentinel */
2382};
2383
Christian Heimesc3f30c42008-02-22 16:37:40 +00002384PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002385"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002386\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002387Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002388For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2389The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2390cycle in a manner similar to an odometer (with the rightmost element changing\n\
2391on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002392To compute the product of an iterable with itself, specify the number\n\
2393of repetitions with the optional repeat keyword argument. For example,\n\
2394product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002395product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2396product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2397
2398static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002399 PyVarObject_HEAD_INIT(NULL, 0)
2400 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002401 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002402 0, /* tp_itemsize */
2403 /* methods */
2404 (destructor)product_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002405 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002406 0, /* tp_getattr */
2407 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002408 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002409 0, /* tp_repr */
2410 0, /* tp_as_number */
2411 0, /* tp_as_sequence */
2412 0, /* tp_as_mapping */
2413 0, /* tp_hash */
2414 0, /* tp_call */
2415 0, /* tp_str */
2416 PyObject_GenericGetAttr, /* tp_getattro */
2417 0, /* tp_setattro */
2418 0, /* tp_as_buffer */
2419 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2420 Py_TPFLAGS_BASETYPE, /* tp_flags */
2421 product_doc, /* tp_doc */
2422 (traverseproc)product_traverse, /* tp_traverse */
2423 0, /* tp_clear */
2424 0, /* tp_richcompare */
2425 0, /* tp_weaklistoffset */
2426 PyObject_SelfIter, /* tp_iter */
2427 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002428 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002429 0, /* tp_members */
2430 0, /* tp_getset */
2431 0, /* tp_base */
2432 0, /* tp_dict */
2433 0, /* tp_descr_get */
2434 0, /* tp_descr_set */
2435 0, /* tp_dictoffset */
2436 0, /* tp_init */
2437 0, /* tp_alloc */
2438 product_new, /* tp_new */
2439 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002440};
2441
2442
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002443/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002444
2445typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002446 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002447 PyObject *pool; /* input converted to a tuple */
2448 Py_ssize_t *indices; /* one index per result element */
2449 PyObject *result; /* most recently returned result tuple */
2450 Py_ssize_t r; /* size of result tuple */
2451 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002452} combinationsobject;
2453
2454static PyTypeObject combinations_type;
2455
Tal Einatc4bccd32018-09-12 00:49:13 +03002456
2457/*[clinic input]
2458@classmethod
2459itertools.combinations.__new__
2460 iterable: object
2461 r: Py_ssize_t
2462Return successive r-length combinations of elements in the iterable.
2463
2464combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)
2465[clinic start generated code]*/
2466
Christian Heimes380f7f22008-02-28 11:19:05 +00002467static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002468itertools_combinations_impl(PyTypeObject *type, PyObject *iterable,
2469 Py_ssize_t r)
2470/*[clinic end generated code: output=87a689b39c40039c input=06bede09e3da20f8]*/
Christian Heimes380f7f22008-02-28 11:19:05 +00002471{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002472 combinationsobject *co;
2473 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002474 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002475 Py_ssize_t *indices = NULL;
2476 Py_ssize_t i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002477
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002478 pool = PySequence_Tuple(iterable);
2479 if (pool == NULL)
2480 goto error;
2481 n = PyTuple_GET_SIZE(pool);
2482 if (r < 0) {
2483 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2484 goto error;
2485 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002486
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002487 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002488 if (indices == NULL) {
2489 PyErr_NoMemory();
2490 goto error;
2491 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002492
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002493 for (i=0 ; i<r ; i++)
2494 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002495
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002496 /* create combinationsobject structure */
2497 co = (combinationsobject *)type->tp_alloc(type, 0);
2498 if (co == NULL)
2499 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002500
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002501 co->pool = pool;
2502 co->indices = indices;
2503 co->result = NULL;
2504 co->r = r;
2505 co->stopped = r > n ? 1 : 0;
2506
2507 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002508
2509error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002510 if (indices != NULL)
2511 PyMem_Free(indices);
2512 Py_XDECREF(pool);
2513 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002514}
2515
2516static void
2517combinations_dealloc(combinationsobject *co)
2518{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002519 PyObject_GC_UnTrack(co);
2520 Py_XDECREF(co->pool);
2521 Py_XDECREF(co->result);
2522 if (co->indices != NULL)
2523 PyMem_Free(co->indices);
2524 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002525}
2526
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002527static PyObject *
2528combinations_sizeof(combinationsobject *co, void *unused)
2529{
2530 Py_ssize_t res;
2531
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002532 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002533 res += co->r * sizeof(Py_ssize_t);
2534 return PyLong_FromSsize_t(res);
2535}
2536
Christian Heimes380f7f22008-02-28 11:19:05 +00002537static int
2538combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2539{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002540 Py_VISIT(co->pool);
2541 Py_VISIT(co->result);
2542 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002543}
2544
2545static PyObject *
2546combinations_next(combinationsobject *co)
2547{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002548 PyObject *elem;
2549 PyObject *oldelem;
2550 PyObject *pool = co->pool;
2551 Py_ssize_t *indices = co->indices;
2552 PyObject *result = co->result;
2553 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2554 Py_ssize_t r = co->r;
2555 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002556
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002557 if (co->stopped)
2558 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002559
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002560 if (result == NULL) {
2561 /* On the first pass, initialize result tuple using the indices */
2562 result = PyTuple_New(r);
2563 if (result == NULL)
2564 goto empty;
2565 co->result = result;
2566 for (i=0; i<r ; i++) {
2567 index = indices[i];
2568 elem = PyTuple_GET_ITEM(pool, index);
2569 Py_INCREF(elem);
2570 PyTuple_SET_ITEM(result, i, elem);
2571 }
2572 } else {
2573 /* Copy the previous result tuple or re-use it if available */
2574 if (Py_REFCNT(result) > 1) {
2575 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002576 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002577 if (result == NULL)
2578 goto empty;
2579 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002580 Py_DECREF(old_result);
2581 }
2582 /* Now, we've got the only copy so we can update it in-place
2583 * CPython's empty tuple is a singleton and cached in
2584 * PyTuple's freelist.
2585 */
2586 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002587
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002588 /* Scan indices right-to-left until finding one that is not
2589 at its maximum (i + n - r). */
2590 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2591 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002592
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002593 /* If i is negative, then the indices are all at
2594 their maximum value and we're done. */
2595 if (i < 0)
2596 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002597
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002598 /* Increment the current index which we know is not at its
2599 maximum. Then move back to the right setting each index
2600 to its lowest possible value (one higher than the index
2601 to its left -- this maintains the sort order invariant). */
2602 indices[i]++;
2603 for (j=i+1 ; j<r ; j++)
2604 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002605
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002606 /* Update the result tuple for the new indices
2607 starting with i, the leftmost index that changed */
2608 for ( ; i<r ; i++) {
2609 index = indices[i];
2610 elem = PyTuple_GET_ITEM(pool, index);
2611 Py_INCREF(elem);
2612 oldelem = PyTuple_GET_ITEM(result, i);
2613 PyTuple_SET_ITEM(result, i, elem);
2614 Py_DECREF(oldelem);
2615 }
2616 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002617
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002618 Py_INCREF(result);
2619 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002620
2621empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002622 co->stopped = 1;
2623 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002624}
2625
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002626static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302627combinations_reduce(combinationsobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002628{
2629 if (lz->result == NULL) {
2630 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2631 } else if (lz->stopped) {
2632 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2633 } else {
2634 PyObject *indices;
2635 Py_ssize_t i;
2636
2637 /* we must pickle the indices and use them for setstate */
2638 indices = PyTuple_New(lz->r);
2639 if (!indices)
2640 return NULL;
2641 for (i=0; i<lz->r; i++)
2642 {
2643 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2644 if (!index) {
2645 Py_DECREF(indices);
2646 return NULL;
2647 }
2648 PyTuple_SET_ITEM(indices, i, index);
2649 }
2650
2651 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2652 }
2653}
2654
2655static PyObject *
2656combinations_setstate(combinationsobject *lz, PyObject *state)
2657{
2658 PyObject *result;
2659 Py_ssize_t i;
2660 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2661
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002662 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002663 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2664 return NULL;
2665 }
2666
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002667 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002668 Py_ssize_t max;
2669 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2670 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002671
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002672 if (index == -1 && PyErr_Occurred())
2673 return NULL; /* not an integer */
2674 max = i + n - lz->r;
2675 /* clamp the index (beware of negative max) */
2676 if (index > max)
2677 index = max;
2678 if (index < 0)
2679 index = 0;
2680 lz->indices[i] = index;
2681 }
2682
2683 result = PyTuple_New(lz->r);
2684 if (result == NULL)
2685 return NULL;
2686 for (i=0; i<lz->r; i++) {
2687 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2688 Py_INCREF(element);
2689 PyTuple_SET_ITEM(result, i, element);
2690 }
2691
Serhiy Storchaka48842712016-04-06 09:45:48 +03002692 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002693 Py_RETURN_NONE;
2694}
2695
2696static PyMethodDef combinations_methods[] = {
2697 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2698 reduce_doc},
2699 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2700 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002701 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2702 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002703 {NULL, NULL} /* sentinel */
2704};
2705
Christian Heimes380f7f22008-02-28 11:19:05 +00002706static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002707 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002708 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002709 sizeof(combinationsobject), /* tp_basicsize */
2710 0, /* tp_itemsize */
2711 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002712 (destructor)combinations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002713 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002714 0, /* tp_getattr */
2715 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002716 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002717 0, /* tp_repr */
2718 0, /* tp_as_number */
2719 0, /* tp_as_sequence */
2720 0, /* tp_as_mapping */
2721 0, /* tp_hash */
2722 0, /* tp_call */
2723 0, /* tp_str */
2724 PyObject_GenericGetAttr, /* tp_getattro */
2725 0, /* tp_setattro */
2726 0, /* tp_as_buffer */
2727 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2728 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03002729 itertools_combinations__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002730 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002731 0, /* tp_clear */
2732 0, /* tp_richcompare */
2733 0, /* tp_weaklistoffset */
2734 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002735 (iternextfunc)combinations_next, /* tp_iternext */
2736 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002737 0, /* tp_members */
2738 0, /* tp_getset */
2739 0, /* tp_base */
2740 0, /* tp_dict */
2741 0, /* tp_descr_get */
2742 0, /* tp_descr_set */
2743 0, /* tp_dictoffset */
2744 0, /* tp_init */
2745 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03002746 itertools_combinations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002747 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002748};
2749
2750
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002751/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002752
2753/* Equivalent to:
2754
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002755 def combinations_with_replacement(iterable, r):
2756 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2757 # number items returned: (n+r-1)! / r! / (n-1)!
2758 pool = tuple(iterable)
2759 n = len(pool)
2760 indices = [0] * r
2761 yield tuple(pool[i] for i in indices)
2762 while 1:
2763 for i in reversed(range(r)):
2764 if indices[i] != n - 1:
2765 break
2766 else:
2767 return
2768 indices[i:] = [indices[i] + 1] * (r - i)
2769 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002770
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002771 def combinations_with_replacement2(iterable, r):
2772 'Alternate version that filters from product()'
2773 pool = tuple(iterable)
2774 n = len(pool)
2775 for indices in product(range(n), repeat=r):
2776 if sorted(indices) == list(indices):
2777 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002778*/
2779typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002780 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002781 PyObject *pool; /* input converted to a tuple */
2782 Py_ssize_t *indices; /* one index per result element */
2783 PyObject *result; /* most recently returned result tuple */
2784 Py_ssize_t r; /* size of result tuple */
2785 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002786} cwrobject;
2787
2788static PyTypeObject cwr_type;
2789
Tal Einatc4bccd32018-09-12 00:49:13 +03002790/*[clinic input]
2791@classmethod
2792itertools.combinations_with_replacement.__new__
2793 iterable: object
2794 r: Py_ssize_t
2795Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats.
2796
2797combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2798[clinic start generated code]*/
2799
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002800static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002801itertools_combinations_with_replacement_impl(PyTypeObject *type,
2802 PyObject *iterable,
2803 Py_ssize_t r)
2804/*[clinic end generated code: output=48b26856d4e659ca input=dc2a8c7ba785fad7]*/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002805{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002806 cwrobject *co;
2807 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002808 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002809 Py_ssize_t *indices = NULL;
2810 Py_ssize_t i;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002811
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002812 pool = PySequence_Tuple(iterable);
2813 if (pool == NULL)
2814 goto error;
2815 n = PyTuple_GET_SIZE(pool);
2816 if (r < 0) {
2817 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2818 goto error;
2819 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002820
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002821 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002822 if (indices == NULL) {
2823 PyErr_NoMemory();
2824 goto error;
2825 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002826
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002827 for (i=0 ; i<r ; i++)
2828 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002829
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002830 /* create cwrobject structure */
2831 co = (cwrobject *)type->tp_alloc(type, 0);
2832 if (co == NULL)
2833 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002834
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002835 co->pool = pool;
2836 co->indices = indices;
2837 co->result = NULL;
2838 co->r = r;
2839 co->stopped = !n && r;
2840
2841 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002842
2843error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002844 if (indices != NULL)
2845 PyMem_Free(indices);
2846 Py_XDECREF(pool);
2847 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002848}
2849
2850static void
2851cwr_dealloc(cwrobject *co)
2852{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002853 PyObject_GC_UnTrack(co);
2854 Py_XDECREF(co->pool);
2855 Py_XDECREF(co->result);
2856 if (co->indices != NULL)
2857 PyMem_Free(co->indices);
2858 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002859}
2860
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002861static PyObject *
2862cwr_sizeof(cwrobject *co, void *unused)
2863{
2864 Py_ssize_t res;
2865
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002866 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002867 res += co->r * sizeof(Py_ssize_t);
2868 return PyLong_FromSsize_t(res);
2869}
2870
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002871static int
2872cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2873{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002874 Py_VISIT(co->pool);
2875 Py_VISIT(co->result);
2876 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002877}
2878
2879static PyObject *
2880cwr_next(cwrobject *co)
2881{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002882 PyObject *elem;
2883 PyObject *oldelem;
2884 PyObject *pool = co->pool;
2885 Py_ssize_t *indices = co->indices;
2886 PyObject *result = co->result;
2887 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2888 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002889 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002890
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002891 if (co->stopped)
2892 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002893
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002894 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002895 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002896 result = PyTuple_New(r);
2897 if (result == NULL)
2898 goto empty;
2899 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002900 if (n > 0) {
2901 elem = PyTuple_GET_ITEM(pool, 0);
2902 for (i=0; i<r ; i++) {
2903 assert(indices[i] == 0);
2904 Py_INCREF(elem);
2905 PyTuple_SET_ITEM(result, i, elem);
2906 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002907 }
2908 } else {
2909 /* Copy the previous result tuple or re-use it if available */
2910 if (Py_REFCNT(result) > 1) {
2911 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002912 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002913 if (result == NULL)
2914 goto empty;
2915 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002916 Py_DECREF(old_result);
2917 }
2918 /* Now, we've got the only copy so we can update it in-place CPython's
2919 empty tuple is a singleton and cached in PyTuple's freelist. */
2920 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002921
Tim Peters9edb1682013-09-03 11:49:31 -05002922 /* Scan indices right-to-left until finding one that is not
2923 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002924 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2925 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002926
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002927 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002928 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002929 if (i < 0)
2930 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002931
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002932 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002933 maximum. Then set all to the right to the same value. */
2934 index = indices[i] + 1;
2935 assert(index < n);
2936 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002937 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002938 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002939 Py_INCREF(elem);
2940 oldelem = PyTuple_GET_ITEM(result, i);
2941 PyTuple_SET_ITEM(result, i, elem);
2942 Py_DECREF(oldelem);
2943 }
2944 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002945
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002946 Py_INCREF(result);
2947 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002948
2949empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002950 co->stopped = 1;
2951 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002952}
2953
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002954static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302955cwr_reduce(cwrobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002956{
2957 if (lz->result == NULL) {
2958 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2959 } else if (lz->stopped) {
2960 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2961 } else {
2962 PyObject *indices;
2963 Py_ssize_t i;
2964
2965 /* we must pickle the indices and use them for setstate */
2966 indices = PyTuple_New(lz->r);
2967 if (!indices)
2968 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002969 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002970 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2971 if (!index) {
2972 Py_DECREF(indices);
2973 return NULL;
2974 }
2975 PyTuple_SET_ITEM(indices, i, index);
2976 }
2977
2978 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2979 }
2980}
2981
2982static PyObject *
2983cwr_setstate(cwrobject *lz, PyObject *state)
2984{
2985 PyObject *result;
2986 Py_ssize_t n, i;
2987
2988 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2989 {
2990 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2991 return NULL;
2992 }
2993
2994 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002995 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002996 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2997 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002998
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002999 if (index < 0 && PyErr_Occurred())
3000 return NULL; /* not an integer */
3001 /* clamp the index */
3002 if (index < 0)
3003 index = 0;
3004 else if (index > n-1)
3005 index = n-1;
3006 lz->indices[i] = index;
3007 }
3008 result = PyTuple_New(lz->r);
3009 if (result == NULL)
3010 return NULL;
3011 for (i=0; i<lz->r; i++) {
3012 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
3013 Py_INCREF(element);
3014 PyTuple_SET_ITEM(result, i, element);
3015 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003016 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003017 Py_RETURN_NONE;
3018}
3019
3020static PyMethodDef cwr_methods[] = {
3021 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
3022 reduce_doc},
3023 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
3024 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003025 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
3026 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003027 {NULL, NULL} /* sentinel */
3028};
3029
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003030static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003031 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003032 "itertools.combinations_with_replacement", /* tp_name */
3033 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003034 0, /* tp_itemsize */
3035 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003036 (destructor)cwr_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003037 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003038 0, /* tp_getattr */
3039 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003040 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003041 0, /* tp_repr */
3042 0, /* tp_as_number */
3043 0, /* tp_as_sequence */
3044 0, /* tp_as_mapping */
3045 0, /* tp_hash */
3046 0, /* tp_call */
3047 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003048 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003049 0, /* tp_setattro */
3050 0, /* tp_as_buffer */
3051 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003052 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003053 itertools_combinations_with_replacement__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003054 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003055 0, /* tp_clear */
3056 0, /* tp_richcompare */
3057 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003058 PyObject_SelfIter, /* tp_iter */
3059 (iternextfunc)cwr_next, /* tp_iternext */
3060 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003061 0, /* tp_members */
3062 0, /* tp_getset */
3063 0, /* tp_base */
3064 0, /* tp_dict */
3065 0, /* tp_descr_get */
3066 0, /* tp_descr_set */
3067 0, /* tp_dictoffset */
3068 0, /* tp_init */
3069 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003070 itertools_combinations_with_replacement, /* tp_new */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003071 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003072};
3073
3074
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003075/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003076
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003077def permutations(iterable, r=None):
3078 'permutations(range(3), 2) --> (0,1) (0,2) (1,0) (1,2) (2,0) (2,1)'
3079 pool = tuple(iterable)
3080 n = len(pool)
3081 r = n if r is None else r
3082 indices = range(n)
3083 cycles = range(n-r+1, n+1)[::-1]
3084 yield tuple(pool[i] for i in indices[:r])
3085 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003086 for i in reversed(range(r)):
3087 cycles[i] -= 1
3088 if cycles[i] == 0:
3089 indices[i:] = indices[i+1:] + indices[i:i+1]
3090 cycles[i] = n - i
3091 else:
3092 j = cycles[i]
3093 indices[i], indices[-j] = indices[-j], indices[i]
3094 yield tuple(pool[i] for i in indices[:r])
3095 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003096 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003097 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003098*/
3099
3100typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003101 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003102 PyObject *pool; /* input converted to a tuple */
3103 Py_ssize_t *indices; /* one index per element in the pool */
3104 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3105 PyObject *result; /* most recently returned result tuple */
3106 Py_ssize_t r; /* size of result tuple */
3107 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003108} permutationsobject;
3109
3110static PyTypeObject permutations_type;
3111
Tal Einatc4bccd32018-09-12 00:49:13 +03003112/*[clinic input]
3113@classmethod
3114itertools.permutations.__new__
3115 iterable: object
3116 r as robj: object = None
3117Return successive r-length permutations of elements in the iterable.
3118
3119permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)
3120[clinic start generated code]*/
3121
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003122static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003123itertools_permutations_impl(PyTypeObject *type, PyObject *iterable,
3124 PyObject *robj)
3125/*[clinic end generated code: output=296a72fa76d620ea input=57d0170a4ac0ec7a]*/
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003126{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003127 permutationsobject *po;
3128 Py_ssize_t n;
3129 Py_ssize_t r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003130 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003131 Py_ssize_t *indices = NULL;
3132 Py_ssize_t *cycles = NULL;
3133 Py_ssize_t i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003134
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003135 pool = PySequence_Tuple(iterable);
3136 if (pool == NULL)
3137 goto error;
3138 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003139
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003140 r = n;
3141 if (robj != Py_None) {
3142 if (!PyLong_Check(robj)) {
3143 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3144 goto error;
3145 }
3146 r = PyLong_AsSsize_t(robj);
3147 if (r == -1 && PyErr_Occurred())
3148 goto error;
3149 }
3150 if (r < 0) {
3151 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3152 goto error;
3153 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003154
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003155 indices = PyMem_New(Py_ssize_t, n);
3156 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003157 if (indices == NULL || cycles == NULL) {
3158 PyErr_NoMemory();
3159 goto error;
3160 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003161
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003162 for (i=0 ; i<n ; i++)
3163 indices[i] = i;
3164 for (i=0 ; i<r ; i++)
3165 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003166
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003167 /* create permutationsobject structure */
3168 po = (permutationsobject *)type->tp_alloc(type, 0);
3169 if (po == NULL)
3170 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003171
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003172 po->pool = pool;
3173 po->indices = indices;
3174 po->cycles = cycles;
3175 po->result = NULL;
3176 po->r = r;
3177 po->stopped = r > n ? 1 : 0;
3178
3179 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003180
3181error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003182 if (indices != NULL)
3183 PyMem_Free(indices);
3184 if (cycles != NULL)
3185 PyMem_Free(cycles);
3186 Py_XDECREF(pool);
3187 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003188}
3189
3190static void
3191permutations_dealloc(permutationsobject *po)
3192{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003193 PyObject_GC_UnTrack(po);
3194 Py_XDECREF(po->pool);
3195 Py_XDECREF(po->result);
3196 PyMem_Free(po->indices);
3197 PyMem_Free(po->cycles);
3198 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003199}
3200
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003201static PyObject *
3202permutations_sizeof(permutationsobject *po, void *unused)
3203{
3204 Py_ssize_t res;
3205
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003206 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003207 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3208 res += po->r * sizeof(Py_ssize_t);
3209 return PyLong_FromSsize_t(res);
3210}
3211
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003212static int
3213permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3214{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003215 Py_VISIT(po->pool);
3216 Py_VISIT(po->result);
3217 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003218}
3219
3220static PyObject *
3221permutations_next(permutationsobject *po)
3222{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003223 PyObject *elem;
3224 PyObject *oldelem;
3225 PyObject *pool = po->pool;
3226 Py_ssize_t *indices = po->indices;
3227 Py_ssize_t *cycles = po->cycles;
3228 PyObject *result = po->result;
3229 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3230 Py_ssize_t r = po->r;
3231 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003232
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003233 if (po->stopped)
3234 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003235
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003236 if (result == NULL) {
3237 /* On the first pass, initialize result tuple using the indices */
3238 result = PyTuple_New(r);
3239 if (result == NULL)
3240 goto empty;
3241 po->result = result;
3242 for (i=0; i<r ; i++) {
3243 index = indices[i];
3244 elem = PyTuple_GET_ITEM(pool, index);
3245 Py_INCREF(elem);
3246 PyTuple_SET_ITEM(result, i, elem);
3247 }
3248 } else {
3249 if (n == 0)
3250 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003251
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003252 /* Copy the previous result tuple or re-use it if available */
3253 if (Py_REFCNT(result) > 1) {
3254 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05003255 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003256 if (result == NULL)
3257 goto empty;
3258 po->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003259 Py_DECREF(old_result);
3260 }
3261 /* Now, we've got the only copy so we can update it in-place */
3262 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003263
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003264 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3265 for (i=r-1 ; i>=0 ; i--) {
3266 cycles[i] -= 1;
3267 if (cycles[i] == 0) {
3268 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3269 index = indices[i];
3270 for (j=i ; j<n-1 ; j++)
3271 indices[j] = indices[j+1];
3272 indices[n-1] = index;
3273 cycles[i] = n - i;
3274 } else {
3275 j = cycles[i];
3276 index = indices[i];
3277 indices[i] = indices[n-j];
3278 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003279
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003280 for (k=i; k<r ; k++) {
3281 /* start with i, the leftmost element that changed */
3282 /* yield tuple(pool[k] for k in indices[:r]) */
3283 index = indices[k];
3284 elem = PyTuple_GET_ITEM(pool, index);
3285 Py_INCREF(elem);
3286 oldelem = PyTuple_GET_ITEM(result, k);
3287 PyTuple_SET_ITEM(result, k, elem);
3288 Py_DECREF(oldelem);
3289 }
3290 break;
3291 }
3292 }
3293 /* If i is negative, then the cycles have all
3294 rolled-over and we're done. */
3295 if (i < 0)
3296 goto empty;
3297 }
3298 Py_INCREF(result);
3299 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003300
3301empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003302 po->stopped = 1;
3303 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003304}
3305
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003306static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303307permutations_reduce(permutationsobject *po, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003308{
3309 if (po->result == NULL) {
3310 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3311 } else if (po->stopped) {
3312 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3313 } else {
3314 PyObject *indices=NULL, *cycles=NULL;
3315 Py_ssize_t n, i;
3316
3317 /* we must pickle the indices and cycles and use them for setstate */
3318 n = PyTuple_GET_SIZE(po->pool);
3319 indices = PyTuple_New(n);
3320 if (indices == NULL)
3321 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003322 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003323 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3324 if (!index)
3325 goto err;
3326 PyTuple_SET_ITEM(indices, i, index);
3327 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003328
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003329 cycles = PyTuple_New(po->r);
3330 if (cycles == NULL)
3331 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003332 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003333 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3334 if (!index)
3335 goto err;
3336 PyTuple_SET_ITEM(cycles, i, index);
3337 }
3338 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3339 po->pool, po->r,
3340 indices, cycles);
3341 err:
3342 Py_XDECREF(indices);
3343 Py_XDECREF(cycles);
3344 return NULL;
3345 }
3346}
3347
3348static PyObject *
3349permutations_setstate(permutationsobject *po, PyObject *state)
3350{
3351 PyObject *indices, *cycles, *result;
3352 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003353
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003354 if (!PyTuple_Check(state)) {
3355 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3356 return NULL;
3357 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003358 if (!PyArg_ParseTuple(state, "O!O!",
3359 &PyTuple_Type, &indices,
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003360 &PyTuple_Type, &cycles)) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003361 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003362 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003363
3364 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003365 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003366 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3367 return NULL;
3368 }
3369
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003370 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003371 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3372 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3373 if (index < 0 && PyErr_Occurred())
3374 return NULL; /* not an integer */
3375 /* clamp the index */
3376 if (index < 0)
3377 index = 0;
3378 else if (index > n-1)
3379 index = n-1;
3380 po->indices[i] = index;
3381 }
3382
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003383 for (i=0; i<po->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003384 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3385 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3386 if (index < 0 && PyErr_Occurred())
3387 return NULL; /* not an integer */
3388 if (index < 1)
3389 index = 1;
3390 else if (index > n-i)
3391 index = n-i;
3392 po->cycles[i] = index;
3393 }
3394 result = PyTuple_New(po->r);
3395 if (result == NULL)
3396 return NULL;
3397 for (i=0; i<po->r; i++) {
3398 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3399 Py_INCREF(element);
3400 PyTuple_SET_ITEM(result, i, element);
3401 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003402 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003403 Py_RETURN_NONE;
3404}
3405
3406static PyMethodDef permuations_methods[] = {
3407 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3408 reduce_doc},
3409 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3410 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003411 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3412 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003413 {NULL, NULL} /* sentinel */
3414};
3415
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003416static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003417 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003418 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003419 sizeof(permutationsobject), /* tp_basicsize */
3420 0, /* tp_itemsize */
3421 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003422 (destructor)permutations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003423 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003424 0, /* tp_getattr */
3425 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003426 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003427 0, /* tp_repr */
3428 0, /* tp_as_number */
3429 0, /* tp_as_sequence */
3430 0, /* tp_as_mapping */
3431 0, /* tp_hash */
3432 0, /* tp_call */
3433 0, /* tp_str */
3434 PyObject_GenericGetAttr, /* tp_getattro */
3435 0, /* tp_setattro */
3436 0, /* tp_as_buffer */
3437 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3438 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003439 itertools_permutations__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003440 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003441 0, /* tp_clear */
3442 0, /* tp_richcompare */
3443 0, /* tp_weaklistoffset */
3444 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003445 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003446 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003447 0, /* tp_members */
3448 0, /* tp_getset */
3449 0, /* tp_base */
3450 0, /* tp_dict */
3451 0, /* tp_descr_get */
3452 0, /* tp_descr_set */
3453 0, /* tp_dictoffset */
3454 0, /* tp_init */
3455 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003456 itertools_permutations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003457 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003458};
3459
Tal Einatc4bccd32018-09-12 00:49:13 +03003460
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003461/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003462
3463typedef struct {
3464 PyObject_HEAD
3465 PyObject *total;
3466 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003467 PyObject *binop;
Lisa Roach9718b592018-09-23 17:34:59 -07003468 PyObject *initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003469} accumulateobject;
3470
3471static PyTypeObject accumulate_type;
3472
Tal Einatc4bccd32018-09-12 00:49:13 +03003473/*[clinic input]
3474@classmethod
3475itertools.accumulate.__new__
3476 iterable: object
3477 func as binop: object = None
Lisa Roach9718b592018-09-23 17:34:59 -07003478 *
3479 initial: object = None
Tal Einatc4bccd32018-09-12 00:49:13 +03003480Return series of accumulated sums (or other binary function results).
3481[clinic start generated code]*/
3482
Raymond Hettinger482ba772010-12-01 22:48:00 +00003483static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003484itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
Lisa Roach9718b592018-09-23 17:34:59 -07003485 PyObject *binop, PyObject *initial)
3486/*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003487{
Raymond Hettinger482ba772010-12-01 22:48:00 +00003488 PyObject *it;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003489 accumulateobject *lz;
3490
Raymond Hettinger482ba772010-12-01 22:48:00 +00003491 /* Get iterator. */
3492 it = PyObject_GetIter(iterable);
3493 if (it == NULL)
3494 return NULL;
3495
Raymond Hettinger482ba772010-12-01 22:48:00 +00003496 /* create accumulateobject structure */
3497 lz = (accumulateobject *)type->tp_alloc(type, 0);
3498 if (lz == NULL) {
3499 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003500 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003501 }
3502
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003503 if (binop != Py_None) {
3504 Py_XINCREF(binop);
3505 lz->binop = binop;
3506 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003507 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003508 lz->it = it;
Lisa Roach9718b592018-09-23 17:34:59 -07003509 Py_XINCREF(initial);
3510 lz->initial = initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003511 return (PyObject *)lz;
3512}
3513
3514static void
3515accumulate_dealloc(accumulateobject *lz)
3516{
3517 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003518 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003519 Py_XDECREF(lz->total);
3520 Py_XDECREF(lz->it);
Lisa Roach9718b592018-09-23 17:34:59 -07003521 Py_XDECREF(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003522 Py_TYPE(lz)->tp_free(lz);
3523}
3524
3525static int
3526accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3527{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003528 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003529 Py_VISIT(lz->it);
3530 Py_VISIT(lz->total);
Lisa Roach9718b592018-09-23 17:34:59 -07003531 Py_VISIT(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003532 return 0;
3533}
3534
3535static PyObject *
3536accumulate_next(accumulateobject *lz)
3537{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003538 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003539
Lisa Roach9718b592018-09-23 17:34:59 -07003540 if (lz->initial != Py_None) {
3541 lz->total = lz->initial;
3542 Py_INCREF(Py_None);
3543 lz->initial = Py_None;
3544 Py_INCREF(lz->total);
3545 return lz->total;
3546 }
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003547 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003548 if (val == NULL)
3549 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003550
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003551 if (lz->total == NULL) {
3552 Py_INCREF(val);
3553 lz->total = val;
3554 return lz->total;
3555 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003556
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003557 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003558 newtotal = PyNumber_Add(lz->total, val);
3559 else
3560 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003561 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003562 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003563 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003564
Raymond Hettinger482ba772010-12-01 22:48:00 +00003565 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003566 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003567 return newtotal;
3568}
3569
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003570static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303571accumulate_reduce(accumulateobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003572{
Lisa Roach9718b592018-09-23 17:34:59 -07003573 if (lz->initial != Py_None) {
3574 PyObject *it;
3575
3576 assert(lz->total == NULL);
3577 if (PyType_Ready(&chain_type) < 0)
3578 return NULL;
3579 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3580 lz->initial, lz->it);
3581 if (it == NULL)
3582 return NULL;
3583 return Py_BuildValue("O(NO)O", Py_TYPE(lz),
3584 it, lz->binop?lz->binop:Py_None, Py_None);
3585 }
Serhiy Storchakad5516252016-03-06 14:00:45 +02003586 if (lz->total == Py_None) {
3587 PyObject *it;
3588
3589 if (PyType_Ready(&chain_type) < 0)
3590 return NULL;
3591 if (PyType_Ready(&islice_type) < 0)
3592 return NULL;
3593 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3594 lz->total, lz->it);
3595 if (it == NULL)
3596 return NULL;
3597 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3598 it, lz->binop ? lz->binop : Py_None);
3599 if (it == NULL)
3600 return NULL;
3601 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3602 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003603 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3604 lz->it, lz->binop?lz->binop:Py_None,
3605 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003606}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003607
3608static PyObject *
3609accumulate_setstate(accumulateobject *lz, PyObject *state)
3610{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003611 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003612 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003613 Py_RETURN_NONE;
3614}
3615
3616static PyMethodDef accumulate_methods[] = {
3617 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3618 reduce_doc},
3619 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3620 setstate_doc},
3621 {NULL, NULL} /* sentinel */
3622};
3623
Raymond Hettinger482ba772010-12-01 22:48:00 +00003624static PyTypeObject accumulate_type = {
3625 PyVarObject_HEAD_INIT(NULL, 0)
3626 "itertools.accumulate", /* tp_name */
3627 sizeof(accumulateobject), /* tp_basicsize */
3628 0, /* tp_itemsize */
3629 /* methods */
3630 (destructor)accumulate_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003631 0, /* tp_vectorcall_offset */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003632 0, /* tp_getattr */
3633 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003634 0, /* tp_as_async */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003635 0, /* tp_repr */
3636 0, /* tp_as_number */
3637 0, /* tp_as_sequence */
3638 0, /* tp_as_mapping */
3639 0, /* tp_hash */
3640 0, /* tp_call */
3641 0, /* tp_str */
3642 PyObject_GenericGetAttr, /* tp_getattro */
3643 0, /* tp_setattro */
3644 0, /* tp_as_buffer */
3645 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3646 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003647 itertools_accumulate__doc__, /* tp_doc */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003648 (traverseproc)accumulate_traverse, /* tp_traverse */
3649 0, /* tp_clear */
3650 0, /* tp_richcompare */
3651 0, /* tp_weaklistoffset */
3652 PyObject_SelfIter, /* tp_iter */
3653 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003654 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003655 0, /* tp_members */
3656 0, /* tp_getset */
3657 0, /* tp_base */
3658 0, /* tp_dict */
3659 0, /* tp_descr_get */
3660 0, /* tp_descr_set */
3661 0, /* tp_dictoffset */
3662 0, /* tp_init */
3663 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003664 itertools_accumulate, /* tp_new */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003665 PyObject_GC_Del, /* tp_free */
3666};
3667
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003668
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003669/* compress object ************************************************************/
3670
3671/* Equivalent to:
3672
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003673 def compress(data, selectors):
3674 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3675 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003676*/
3677
3678typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003679 PyObject_HEAD
3680 PyObject *data;
3681 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003682} compressobject;
3683
3684static PyTypeObject compress_type;
3685
Tal Einatc4bccd32018-09-12 00:49:13 +03003686/*[clinic input]
3687@classmethod
3688itertools.compress.__new__
3689 data as seq1: object
3690 selectors as seq2: object
3691Return data elements corresponding to true selector elements.
3692
3693Forms a shorter iterator from selected data elements using the selectors to
3694choose the data elements.
3695[clinic start generated code]*/
3696
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003697static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003698itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
3699/*[clinic end generated code: output=7e67157212ed09e0 input=79596d7cd20c77e5]*/
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003700{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003701 PyObject *data=NULL, *selectors=NULL;
3702 compressobject *lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003703
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003704 data = PyObject_GetIter(seq1);
3705 if (data == NULL)
3706 goto fail;
3707 selectors = PyObject_GetIter(seq2);
3708 if (selectors == NULL)
3709 goto fail;
3710
3711 /* create compressobject structure */
3712 lz = (compressobject *)type->tp_alloc(type, 0);
3713 if (lz == NULL)
3714 goto fail;
3715 lz->data = data;
3716 lz->selectors = selectors;
3717 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003718
3719fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003720 Py_XDECREF(data);
3721 Py_XDECREF(selectors);
3722 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003723}
3724
3725static void
3726compress_dealloc(compressobject *lz)
3727{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003728 PyObject_GC_UnTrack(lz);
3729 Py_XDECREF(lz->data);
3730 Py_XDECREF(lz->selectors);
3731 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003732}
3733
3734static int
3735compress_traverse(compressobject *lz, visitproc visit, void *arg)
3736{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003737 Py_VISIT(lz->data);
3738 Py_VISIT(lz->selectors);
3739 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003740}
3741
3742static PyObject *
3743compress_next(compressobject *lz)
3744{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003745 PyObject *data = lz->data, *selectors = lz->selectors;
3746 PyObject *datum, *selector;
3747 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3748 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3749 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003750
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003751 while (1) {
3752 /* Steps: get datum, get selector, evaluate selector.
3753 Order is important (to match the pure python version
3754 in terms of which input gets a chance to raise an
3755 exception first).
3756 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003757
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003758 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003759 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003760 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003761
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003762 selector = selectornext(selectors);
3763 if (selector == NULL) {
3764 Py_DECREF(datum);
3765 return NULL;
3766 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003767
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003768 ok = PyObject_IsTrue(selector);
3769 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003770 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003771 return datum;
3772 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003773 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003774 return NULL;
3775 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003776}
3777
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003778static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303779compress_reduce(compressobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003780{
3781 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3782 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003783}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003784
3785static PyMethodDef compress_methods[] = {
3786 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3787 reduce_doc},
3788 {NULL, NULL} /* sentinel */
3789};
3790
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003791static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003792 PyVarObject_HEAD_INIT(NULL, 0)
3793 "itertools.compress", /* tp_name */
3794 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003795 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003796 /* methods */
3797 (destructor)compress_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003798 0, /* tp_vectorcall_offset */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003799 0, /* tp_getattr */
3800 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003801 0, /* tp_as_async */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003802 0, /* tp_repr */
3803 0, /* tp_as_number */
3804 0, /* tp_as_sequence */
3805 0, /* tp_as_mapping */
3806 0, /* tp_hash */
3807 0, /* tp_call */
3808 0, /* tp_str */
3809 PyObject_GenericGetAttr, /* tp_getattro */
3810 0, /* tp_setattro */
3811 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003812 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003813 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003814 itertools_compress__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003815 (traverseproc)compress_traverse, /* tp_traverse */
3816 0, /* tp_clear */
3817 0, /* tp_richcompare */
3818 0, /* tp_weaklistoffset */
3819 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003820 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003821 compress_methods, /* tp_methods */
3822 0, /* tp_members */
3823 0, /* tp_getset */
3824 0, /* tp_base */
3825 0, /* tp_dict */
3826 0, /* tp_descr_get */
3827 0, /* tp_descr_set */
3828 0, /* tp_dictoffset */
3829 0, /* tp_init */
3830 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003831 itertools_compress, /* tp_new */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003832 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003833};
3834
3835
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003836/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003837
3838typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003839 PyObject_HEAD
3840 PyObject *func;
3841 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003842} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003843
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003844static PyTypeObject filterfalse_type;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003845
Tal Einatc4bccd32018-09-12 00:49:13 +03003846/*[clinic input]
3847@classmethod
3848itertools.filterfalse.__new__
3849 function as func: object
3850 iterable as seq: object
3851 /
3852Return those items of iterable for which function(item) is false.
3853
3854If function is None, return the items that are false.
3855[clinic start generated code]*/
3856
Raymond Hettinger60eca932003-02-09 06:40:58 +00003857static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003858itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
3859/*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003860{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003861 PyObject *it;
3862 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003863
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003864 /* Get iterator. */
3865 it = PyObject_GetIter(seq);
3866 if (it == NULL)
3867 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003868
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003869 /* create filterfalseobject structure */
3870 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3871 if (lz == NULL) {
3872 Py_DECREF(it);
3873 return NULL;
3874 }
3875 Py_INCREF(func);
3876 lz->func = func;
3877 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003878
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003879 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003880}
3881
3882static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003883filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003884{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003885 PyObject_GC_UnTrack(lz);
3886 Py_XDECREF(lz->func);
3887 Py_XDECREF(lz->it);
3888 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003889}
3890
3891static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003892filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003893{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003894 Py_VISIT(lz->it);
3895 Py_VISIT(lz->func);
3896 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003897}
3898
3899static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003900filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003901{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003902 PyObject *item;
3903 PyObject *it = lz->it;
3904 long ok;
3905 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003906
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003907 iternext = *Py_TYPE(it)->tp_iternext;
3908 for (;;) {
3909 item = iternext(it);
3910 if (item == NULL)
3911 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003912
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003913 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3914 ok = PyObject_IsTrue(item);
3915 } else {
3916 PyObject *good;
Jeroen Demeyer196a5302019-07-04 12:31:34 +02003917 good = _PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003918 if (good == NULL) {
3919 Py_DECREF(item);
3920 return NULL;
3921 }
3922 ok = PyObject_IsTrue(good);
3923 Py_DECREF(good);
3924 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003925 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003926 return item;
3927 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003928 if (ok < 0)
3929 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003930 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003931}
3932
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003933static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303934filterfalse_reduce(filterfalseobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003935{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003936 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3937}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003938
3939static PyMethodDef filterfalse_methods[] = {
3940 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3941 reduce_doc},
3942 {NULL, NULL} /* sentinel */
3943};
3944
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003945static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003946 PyVarObject_HEAD_INIT(NULL, 0)
3947 "itertools.filterfalse", /* tp_name */
3948 sizeof(filterfalseobject), /* tp_basicsize */
3949 0, /* tp_itemsize */
3950 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003951 (destructor)filterfalse_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003952 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003953 0, /* tp_getattr */
3954 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003955 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003956 0, /* tp_repr */
3957 0, /* tp_as_number */
3958 0, /* tp_as_sequence */
3959 0, /* tp_as_mapping */
3960 0, /* tp_hash */
3961 0, /* tp_call */
3962 0, /* tp_str */
3963 PyObject_GenericGetAttr, /* tp_getattro */
3964 0, /* tp_setattro */
3965 0, /* tp_as_buffer */
3966 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3967 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003968 itertools_filterfalse__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003969 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003970 0, /* tp_clear */
3971 0, /* tp_richcompare */
3972 0, /* tp_weaklistoffset */
3973 PyObject_SelfIter, /* tp_iter */
3974 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003975 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003976 0, /* tp_members */
3977 0, /* tp_getset */
3978 0, /* tp_base */
3979 0, /* tp_dict */
3980 0, /* tp_descr_get */
3981 0, /* tp_descr_set */
3982 0, /* tp_dictoffset */
3983 0, /* tp_init */
3984 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003985 itertools_filterfalse, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003986 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003987};
3988
3989
3990/* count object ************************************************************/
3991
3992typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003993 PyObject_HEAD
3994 Py_ssize_t cnt;
3995 PyObject *long_cnt;
3996 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003997} countobject;
3998
Raymond Hettinger30729212009-02-12 06:28:27 +00003999/* Counting logic and invariants:
4000
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004001fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00004002
Serhiy Storchaka483405b2015-02-17 10:14:30 +02004003 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004004 Advances with: cnt += 1
4005 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00004006
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004007slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00004008
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004009 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
4010 All counting is done with python objects (no overflows or underflows).
4011 Advances with: long_cnt += long_step
4012 Step may be zero -- effectively a slow version of repeat(cnt).
4013 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00004014*/
4015
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004016static PyTypeObject count_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004017
Tal Einatc4bccd32018-09-12 00:49:13 +03004018/*[clinic input]
4019@classmethod
4020itertools.count.__new__
4021 start as long_cnt: object(c_default="NULL") = 0
4022 step as long_step: object(c_default="NULL") = 1
4023Return a count object whose .__next__() method returns consecutive values.
4024
4025Equivalent to:
4026 def count(firstval=0, step=1):
4027 x = firstval
4028 while 1:
4029 yield x
4030 x += step
4031[clinic start generated code]*/
4032
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004033static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03004034itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
4035 PyObject *long_step)
4036/*[clinic end generated code: output=09a9250aebd00b1c input=d7a85eec18bfcd94]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004037{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004038 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004039 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004040 Py_ssize_t cnt = 0;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00004041 long step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004042
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004043 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
4044 (long_step != NULL && !PyNumber_Check(long_step))) {
4045 PyErr_SetString(PyExc_TypeError, "a number is required");
4046 return NULL;
4047 }
Raymond Hettinger30729212009-02-12 06:28:27 +00004048
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004049 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
4050 (long_step == NULL || PyLong_Check(long_step));
4051
4052 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004053 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004054 if (fast_mode) {
4055 assert(PyLong_Check(long_cnt));
4056 cnt = PyLong_AsSsize_t(long_cnt);
4057 if (cnt == -1 && PyErr_Occurred()) {
4058 PyErr_Clear();
4059 fast_mode = 0;
4060 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004061 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004062 } else {
4063 cnt = 0;
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004064 long_cnt = _PyLong_Zero;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004065 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004066 Py_INCREF(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004067
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004068 /* If not specified, step defaults to 1 */
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004069 if (long_step == NULL)
4070 long_step = _PyLong_One;
4071 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004072
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004073 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004074
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004075 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004076 if (fast_mode) {
4077 assert(PyLong_Check(long_step));
4078 step = PyLong_AsLong(long_step);
4079 if (step != 1) {
4080 fast_mode = 0;
4081 if (step == -1 && PyErr_Occurred())
4082 PyErr_Clear();
4083 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004084 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004085
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004086 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004087 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004088 else
4089 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004090
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004091 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4092 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4093 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004094 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004095
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004096 /* create countobject structure */
4097 lz = (countobject *)type->tp_alloc(type, 0);
4098 if (lz == NULL) {
4099 Py_XDECREF(long_cnt);
Zackery Spytz0523c392019-03-26 00:05:29 -06004100 Py_DECREF(long_step);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004101 return NULL;
4102 }
4103 lz->cnt = cnt;
4104 lz->long_cnt = long_cnt;
4105 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004106
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004107 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004108}
4109
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004110static void
4111count_dealloc(countobject *lz)
4112{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004113 PyObject_GC_UnTrack(lz);
4114 Py_XDECREF(lz->long_cnt);
4115 Py_XDECREF(lz->long_step);
4116 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004117}
4118
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004119static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004120count_traverse(countobject *lz, visitproc visit, void *arg)
4121{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004122 Py_VISIT(lz->long_cnt);
4123 Py_VISIT(lz->long_step);
4124 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004125}
4126
4127static PyObject *
4128count_nextlong(countobject *lz)
4129{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004130 PyObject *long_cnt;
4131 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004132
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004133 long_cnt = lz->long_cnt;
4134 if (long_cnt == NULL) {
4135 /* Switch to slow_mode */
4136 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4137 if (long_cnt == NULL)
4138 return NULL;
4139 }
4140 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004141
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004142 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4143 if (stepped_up == NULL)
4144 return NULL;
4145 lz->long_cnt = stepped_up;
4146 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004147}
4148
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004149static PyObject *
4150count_next(countobject *lz)
4151{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004152 if (lz->cnt == PY_SSIZE_T_MAX)
4153 return count_nextlong(lz);
4154 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004155}
4156
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004157static PyObject *
4158count_repr(countobject *lz)
4159{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004160 if (lz->cnt != PY_SSIZE_T_MAX)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004161 return PyUnicode_FromFormat("%s(%zd)",
4162 _PyType_Name(Py_TYPE(lz)), lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004163
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004164 if (PyLong_Check(lz->long_step)) {
4165 long step = PyLong_AsLong(lz->long_step);
4166 if (step == -1 && PyErr_Occurred()) {
4167 PyErr_Clear();
4168 }
4169 if (step == 1) {
4170 /* Don't display step when it is an integer equal to 1 */
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004171 return PyUnicode_FromFormat("%s(%R)",
4172 _PyType_Name(Py_TYPE(lz)),
4173 lz->long_cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004174 }
4175 }
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004176 return PyUnicode_FromFormat("%s(%R, %R)",
4177 _PyType_Name(Py_TYPE(lz)),
4178 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004179}
4180
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004181static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304182count_reduce(countobject *lz, PyObject *Py_UNUSED(ignored))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004183{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004184 if (lz->cnt == PY_SSIZE_T_MAX)
4185 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4186 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004187}
4188
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004189static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004190 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004191 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004192 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004193};
4194
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004195static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004196 PyVarObject_HEAD_INIT(NULL, 0)
4197 "itertools.count", /* tp_name */
4198 sizeof(countobject), /* tp_basicsize */
4199 0, /* tp_itemsize */
4200 /* methods */
4201 (destructor)count_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004202 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004203 0, /* tp_getattr */
4204 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004205 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004206 (reprfunc)count_repr, /* tp_repr */
4207 0, /* tp_as_number */
4208 0, /* tp_as_sequence */
4209 0, /* tp_as_mapping */
4210 0, /* tp_hash */
4211 0, /* tp_call */
4212 0, /* tp_str */
4213 PyObject_GenericGetAttr, /* tp_getattro */
4214 0, /* tp_setattro */
4215 0, /* tp_as_buffer */
4216 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004217 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03004218 itertools_count__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004219 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004220 0, /* tp_clear */
4221 0, /* tp_richcompare */
4222 0, /* tp_weaklistoffset */
4223 PyObject_SelfIter, /* tp_iter */
4224 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004225 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004226 0, /* tp_members */
4227 0, /* tp_getset */
4228 0, /* tp_base */
4229 0, /* tp_dict */
4230 0, /* tp_descr_get */
4231 0, /* tp_descr_set */
4232 0, /* tp_dictoffset */
4233 0, /* tp_init */
4234 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03004235 itertools_count, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004236 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004237};
4238
4239
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004240/* repeat object ************************************************************/
4241
4242typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004243 PyObject_HEAD
4244 PyObject *element;
4245 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004246} repeatobject;
4247
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004248static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004249
4250static PyObject *
4251repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4252{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004253 repeatobject *ro;
4254 PyObject *element;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004255 Py_ssize_t cnt = -1, n_kwds = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004256 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004257
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004258 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4259 &element, &cnt))
4260 return NULL;
Raymond Hettinger7c2bb5b2003-05-03 05:59:48 +00004261
Raymond Hettinger97d35552014-06-24 21:36:58 -07004262 if (kwds != NULL)
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004263 n_kwds = PyDict_GET_SIZE(kwds);
Raymond Hettinger97d35552014-06-24 21:36:58 -07004264 /* Does user supply times argument? */
4265 if ((PyTuple_Size(args) + n_kwds == 2) && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004266 cnt = 0;
4267
4268 ro = (repeatobject *)type->tp_alloc(type, 0);
4269 if (ro == NULL)
4270 return NULL;
4271 Py_INCREF(element);
4272 ro->element = element;
4273 ro->cnt = cnt;
4274 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004275}
4276
4277static void
4278repeat_dealloc(repeatobject *ro)
4279{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004280 PyObject_GC_UnTrack(ro);
4281 Py_XDECREF(ro->element);
4282 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004283}
4284
4285static int
4286repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4287{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004288 Py_VISIT(ro->element);
4289 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004290}
4291
4292static PyObject *
4293repeat_next(repeatobject *ro)
4294{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004295 if (ro->cnt == 0)
4296 return NULL;
4297 if (ro->cnt > 0)
4298 ro->cnt--;
4299 Py_INCREF(ro->element);
4300 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004301}
4302
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004303static PyObject *
4304repeat_repr(repeatobject *ro)
4305{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004306 if (ro->cnt == -1)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004307 return PyUnicode_FromFormat("%s(%R)",
4308 _PyType_Name(Py_TYPE(ro)), ro->element);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004309 else
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004310 return PyUnicode_FromFormat("%s(%R, %zd)",
4311 _PyType_Name(Py_TYPE(ro)), ro->element,
4312 ro->cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004313}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004314
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004315static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304316repeat_len(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004317{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004318 if (ro->cnt == -1) {
4319 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4320 return NULL;
4321 }
4322 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004323}
4324
Armin Rigof5b3e362006-02-11 21:32:43 +00004325PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004326
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004327static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304328repeat_reduce(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004329{
4330 /* unpickle this so that a new repeat iterator is constructed with an
4331 * object, then call __setstate__ on it to set cnt
4332 */
4333 if (ro->cnt >= 0)
4334 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4335 else
4336 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4337}
4338
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004339static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004340 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004341 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004342 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004343};
4344
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004345PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004346"repeat(object [,times]) -> create an iterator which returns the object\n\
4347for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004348endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004349
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004350static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004351 PyVarObject_HEAD_INIT(NULL, 0)
4352 "itertools.repeat", /* tp_name */
4353 sizeof(repeatobject), /* tp_basicsize */
4354 0, /* tp_itemsize */
4355 /* methods */
4356 (destructor)repeat_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004357 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004358 0, /* tp_getattr */
4359 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004360 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004361 (reprfunc)repeat_repr, /* tp_repr */
4362 0, /* tp_as_number */
4363 0, /* tp_as_sequence */
4364 0, /* tp_as_mapping */
4365 0, /* tp_hash */
4366 0, /* tp_call */
4367 0, /* tp_str */
4368 PyObject_GenericGetAttr, /* tp_getattro */
4369 0, /* tp_setattro */
4370 0, /* tp_as_buffer */
4371 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4372 Py_TPFLAGS_BASETYPE, /* tp_flags */
4373 repeat_doc, /* tp_doc */
4374 (traverseproc)repeat_traverse, /* tp_traverse */
4375 0, /* tp_clear */
4376 0, /* tp_richcompare */
4377 0, /* tp_weaklistoffset */
4378 PyObject_SelfIter, /* tp_iter */
4379 (iternextfunc)repeat_next, /* tp_iternext */
4380 repeat_methods, /* tp_methods */
4381 0, /* tp_members */
4382 0, /* tp_getset */
4383 0, /* tp_base */
4384 0, /* tp_dict */
4385 0, /* tp_descr_get */
4386 0, /* tp_descr_set */
4387 0, /* tp_dictoffset */
4388 0, /* tp_init */
4389 0, /* tp_alloc */
4390 repeat_new, /* tp_new */
4391 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004392};
4393
Tal Einatc4bccd32018-09-12 00:49:13 +03004394
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004395/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004396
4397typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004398 PyObject_HEAD
4399 Py_ssize_t tuplesize;
4400 Py_ssize_t numactive;
4401 PyObject *ittuple; /* tuple of iterators */
4402 PyObject *result;
4403 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004404} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004405
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004406static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004407
4408static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004409zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004410{
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004411 _Py_IDENTIFIER(fillvalue);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004412 ziplongestobject *lz;
4413 Py_ssize_t i;
4414 PyObject *ittuple; /* tuple of iterators */
4415 PyObject *result;
4416 PyObject *fillvalue = Py_None;
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004417 Py_ssize_t tuplesize;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004418
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004419 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004420 fillvalue = NULL;
4421 if (PyDict_GET_SIZE(kwds) == 1) {
4422 fillvalue = _PyDict_GetItemIdWithError(kwds, &PyId_fillvalue);
4423 }
4424 if (fillvalue == NULL) {
4425 if (!PyErr_Occurred()) {
4426 PyErr_SetString(PyExc_TypeError,
4427 "zip_longest() got an unexpected keyword argument");
4428 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004429 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004430 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004431 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004432
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004433 /* args must be a tuple */
4434 assert(PyTuple_Check(args));
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004435 tuplesize = PyTuple_GET_SIZE(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004436
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004437 /* obtain iterators */
4438 ittuple = PyTuple_New(tuplesize);
4439 if (ittuple == NULL)
4440 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004441 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004442 PyObject *item = PyTuple_GET_ITEM(args, i);
4443 PyObject *it = PyObject_GetIter(item);
4444 if (it == NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004445 Py_DECREF(ittuple);
4446 return NULL;
4447 }
4448 PyTuple_SET_ITEM(ittuple, i, it);
4449 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004450
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004451 /* create a result holder */
4452 result = PyTuple_New(tuplesize);
4453 if (result == NULL) {
4454 Py_DECREF(ittuple);
4455 return NULL;
4456 }
4457 for (i=0 ; i < tuplesize ; i++) {
4458 Py_INCREF(Py_None);
4459 PyTuple_SET_ITEM(result, i, Py_None);
4460 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004461
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004462 /* create ziplongestobject structure */
4463 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4464 if (lz == NULL) {
4465 Py_DECREF(ittuple);
4466 Py_DECREF(result);
4467 return NULL;
4468 }
4469 lz->ittuple = ittuple;
4470 lz->tuplesize = tuplesize;
4471 lz->numactive = tuplesize;
4472 lz->result = result;
4473 Py_INCREF(fillvalue);
4474 lz->fillvalue = fillvalue;
4475 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004476}
4477
4478static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004479zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004480{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004481 PyObject_GC_UnTrack(lz);
4482 Py_XDECREF(lz->ittuple);
4483 Py_XDECREF(lz->result);
4484 Py_XDECREF(lz->fillvalue);
4485 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004486}
4487
4488static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004489zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004490{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004491 Py_VISIT(lz->ittuple);
4492 Py_VISIT(lz->result);
4493 Py_VISIT(lz->fillvalue);
4494 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004495}
4496
4497static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004498zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004499{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004500 Py_ssize_t i;
4501 Py_ssize_t tuplesize = lz->tuplesize;
4502 PyObject *result = lz->result;
4503 PyObject *it;
4504 PyObject *item;
4505 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004506
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004507 if (tuplesize == 0)
4508 return NULL;
4509 if (lz->numactive == 0)
4510 return NULL;
4511 if (Py_REFCNT(result) == 1) {
4512 Py_INCREF(result);
4513 for (i=0 ; i < tuplesize ; i++) {
4514 it = PyTuple_GET_ITEM(lz->ittuple, i);
4515 if (it == NULL) {
4516 Py_INCREF(lz->fillvalue);
4517 item = lz->fillvalue;
4518 } else {
4519 item = PyIter_Next(it);
4520 if (item == NULL) {
4521 lz->numactive -= 1;
4522 if (lz->numactive == 0 || PyErr_Occurred()) {
4523 lz->numactive = 0;
4524 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004525 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004526 } else {
4527 Py_INCREF(lz->fillvalue);
4528 item = lz->fillvalue;
4529 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4530 Py_DECREF(it);
4531 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004532 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004533 }
4534 olditem = PyTuple_GET_ITEM(result, i);
4535 PyTuple_SET_ITEM(result, i, item);
4536 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004537 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004538 } else {
4539 result = PyTuple_New(tuplesize);
4540 if (result == NULL)
4541 return NULL;
4542 for (i=0 ; i < tuplesize ; i++) {
4543 it = PyTuple_GET_ITEM(lz->ittuple, i);
4544 if (it == NULL) {
4545 Py_INCREF(lz->fillvalue);
4546 item = lz->fillvalue;
4547 } else {
4548 item = PyIter_Next(it);
4549 if (item == NULL) {
4550 lz->numactive -= 1;
4551 if (lz->numactive == 0 || PyErr_Occurred()) {
4552 lz->numactive = 0;
4553 Py_DECREF(result);
4554 return NULL;
4555 } else {
4556 Py_INCREF(lz->fillvalue);
4557 item = lz->fillvalue;
4558 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4559 Py_DECREF(it);
4560 }
4561 }
4562 }
4563 PyTuple_SET_ITEM(result, i, item);
4564 }
4565 }
4566 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004567}
4568
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004569static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304570zip_longest_reduce(ziplongestobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004571{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004572
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004573 /* Create a new tuple with empty sequences where appropriate to pickle.
4574 * Then use setstate to set the fillvalue
4575 */
4576 int i;
4577 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004578
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004579 if (args == NULL)
4580 return NULL;
4581 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4582 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4583 if (elem == NULL) {
4584 elem = PyTuple_New(0);
4585 if (elem == NULL) {
4586 Py_DECREF(args);
4587 return NULL;
4588 }
4589 } else
4590 Py_INCREF(elem);
4591 PyTuple_SET_ITEM(args, i, elem);
4592 }
4593 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4594}
4595
4596static PyObject *
4597zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4598{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004599 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004600 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004601 Py_RETURN_NONE;
4602}
4603
4604static PyMethodDef zip_longest_methods[] = {
4605 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4606 reduce_doc},
4607 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4608 setstate_doc},
4609 {NULL, NULL} /* sentinel */
4610};
4611
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004612PyDoc_STRVAR(zip_longest_doc,
4613"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004614\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004615Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004616the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004617method continues until the longest iterable in the argument sequence\n\
4618is exhausted and then it raises StopIteration. When the shorter iterables\n\
4619are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4620defaults to None or can be specified by a keyword argument.\n\
4621");
4622
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004623static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004624 PyVarObject_HEAD_INIT(NULL, 0)
4625 "itertools.zip_longest", /* tp_name */
4626 sizeof(ziplongestobject), /* tp_basicsize */
4627 0, /* tp_itemsize */
4628 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004629 (destructor)zip_longest_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004630 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004631 0, /* tp_getattr */
4632 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004633 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004634 0, /* tp_repr */
4635 0, /* tp_as_number */
4636 0, /* tp_as_sequence */
4637 0, /* tp_as_mapping */
4638 0, /* tp_hash */
4639 0, /* tp_call */
4640 0, /* tp_str */
4641 PyObject_GenericGetAttr, /* tp_getattro */
4642 0, /* tp_setattro */
4643 0, /* tp_as_buffer */
4644 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4645 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004646 zip_longest_doc, /* tp_doc */
4647 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004648 0, /* tp_clear */
4649 0, /* tp_richcompare */
4650 0, /* tp_weaklistoffset */
4651 PyObject_SelfIter, /* tp_iter */
4652 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004653 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004654 0, /* tp_members */
4655 0, /* tp_getset */
4656 0, /* tp_base */
4657 0, /* tp_dict */
4658 0, /* tp_descr_get */
4659 0, /* tp_descr_set */
4660 0, /* tp_dictoffset */
4661 0, /* tp_init */
4662 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004663 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004664 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004665};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004666
Tal Einatc4bccd32018-09-12 00:49:13 +03004667
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004668/* module level code ********************************************************/
4669
4670PyDoc_STRVAR(module_doc,
4671"Functional tools for creating and using iterators.\n\
4672\n\
4673Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004674count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004675cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004676repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004677\n\
4678Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004679accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
oldkaa0735f2018-02-02 16:52:55 +08004680chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
4681chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004682compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4683dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4684groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004685filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004686islice(seq, [start,] stop [, step]) --> elements from\n\
4687 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004688starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004689tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004690takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
oldkaa0735f2018-02-02 16:52:55 +08004691zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004692\n\
4693Combinatoric generators:\n\
4694product(p, q, ... [repeat=1]) --> cartesian product\n\
4695permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004696combinations(p, r)\n\
4697combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004698");
4699
4700
Raymond Hettingerad983e72003-11-12 14:32:26 +00004701static PyMethodDef module_methods[] = {
Tal Einatc4bccd32018-09-12 00:49:13 +03004702 ITERTOOLS_TEE_METHODDEF
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004703 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +00004704};
4705
Martin v. Löwis1a214512008-06-11 05:26:20 +00004706
4707static struct PyModuleDef itertoolsmodule = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004708 PyModuleDef_HEAD_INIT,
4709 "itertools",
4710 module_doc,
4711 -1,
4712 module_methods,
4713 NULL,
4714 NULL,
4715 NULL,
4716 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00004717};
4718
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004719PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00004720PyInit_itertools(void)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004721{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004722 int i;
4723 PyObject *m;
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004724 const char *name;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004725 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004726 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004727 &combinations_type,
4728 &cwr_type,
4729 &cycle_type,
4730 &dropwhile_type,
4731 &takewhile_type,
4732 &islice_type,
4733 &starmap_type,
4734 &chain_type,
4735 &compress_type,
4736 &filterfalse_type,
4737 &count_type,
4738 &ziplongest_type,
4739 &permutations_type,
4740 &product_type,
4741 &repeat_type,
4742 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004743 &_grouper_type,
4744 &tee_type,
4745 &teedataobject_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004746 NULL
4747 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004748
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004749 Py_TYPE(&teedataobject_type) = &PyType_Type;
4750 m = PyModule_Create(&itertoolsmodule);
4751 if (m == NULL)
4752 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004753
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004754 for (i=0 ; typelist[i] != NULL ; i++) {
4755 if (PyType_Ready(typelist[i]) < 0)
4756 return NULL;
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004757 name = _PyType_Name(typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004758 Py_INCREF(typelist[i]);
Serhiy Storchaka4ab46d72017-09-17 21:11:04 +03004759 PyModule_AddObject(m, name, (PyObject *)typelist[i]);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004760 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004761
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004762 return m;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004763}