blob: ce8b4347ef220bb3ed1a510bc6e88f55f8bf60f0 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
Raymond Hettingerca3788c2015-08-16 14:49:24 -07002#define PY_SSIZE_T_CLEAN
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003#include "Python.h"
Victor Stinner37834132020-10-27 17:12:53 +01004#include "pycore_long.h" // _PyLong_GetZero()
Victor Stinner384621c2020-06-22 17:27:35 +02005#include "pycore_tuple.h" // _PyTuple_ITEMS()
Victor Stinner4a21e572020-04-15 02:35:41 +02006#include <stddef.h> // offsetof()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00007
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00008/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +00009 by Raymond D. Hettinger <python@rcn.com>
Raymond Hettinger96ef8112003-02-01 00:10:11 +000010*/
11
Tal Einat3286ce42018-09-10 21:33:08 +030012/*[clinic input]
13module itertools
14class itertools.groupby "groupbyobject *" "&groupby_type"
15class itertools._grouper "_grouperobject *" "&_grouper_type"
Tal Einatc4bccd32018-09-12 00:49:13 +030016class itertools.teedataobject "teedataobject *" "&teedataobject_type"
17class itertools._tee "teeobject *" "&tee_type"
18class itertools.cycle "cycleobject *" "&cycle_type"
19class itertools.dropwhile "dropwhileobject *" "&dropwhile_type"
20class itertools.takewhile "takewhileobject *" "&takewhile_type"
21class itertools.starmap "starmapobject *" "&starmap_type"
22class itertools.chain "chainobject *" "&chain_type"
23class itertools.combinations "combinationsobject *" "&combinations_type"
24class itertools.combinations_with_replacement "cwr_object *" "&cwr_type"
25class itertools.permutations "permutationsobject *" "&permutations_type"
26class itertools.accumulate "accumulateobject *" "&accumulate_type"
27class itertools.compress "compressobject *" "&compress_type"
28class itertools.filterfalse "filterfalseobject *" "&filterfalse_type"
29class itertools.count "countobject *" "&count_type"
Tal Einat3286ce42018-09-10 21:33:08 +030030[clinic start generated code]*/
Tal Einatc4bccd32018-09-12 00:49:13 +030031/*[clinic end generated code: output=da39a3ee5e6b4b0d input=ea05c93c6d94726a]*/
Tal Einat3286ce42018-09-10 21:33:08 +030032
33static PyTypeObject groupby_type;
34static PyTypeObject _grouper_type;
Tal Einatc4bccd32018-09-12 00:49:13 +030035static PyTypeObject teedataobject_type;
36static PyTypeObject tee_type;
37static PyTypeObject cycle_type;
38static PyTypeObject dropwhile_type;
39static PyTypeObject takewhile_type;
40static PyTypeObject starmap_type;
41static PyTypeObject combinations_type;
42static PyTypeObject cwr_type;
43static PyTypeObject permutations_type;
44static PyTypeObject accumulate_type;
45static PyTypeObject compress_type;
46static PyTypeObject filterfalse_type;
47static PyTypeObject count_type;
48
Tal Einat3286ce42018-09-10 21:33:08 +030049#include "clinic/itertoolsmodule.c.h"
50
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000051
Raymond Hettingera6ea44a2015-08-17 23:55:28 -070052/* groupby object ************************************************************/
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000053
54typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000055 PyObject_HEAD
56 PyObject *it;
57 PyObject *keyfunc;
58 PyObject *tgtkey;
59 PyObject *currkey;
60 PyObject *currvalue;
Serhiy Storchakac247caf2017-09-24 13:36:11 +030061 const void *currgrouper; /* borrowed reference */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000062} groupbyobject;
63
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000064static PyObject *_grouper_create(groupbyobject *, PyObject *);
65
Tal Einat3286ce42018-09-10 21:33:08 +030066/*[clinic input]
67@classmethod
68itertools.groupby.__new__
Raymond Hettingerd25c1c62003-12-06 16:23:06 +000069
Tal Einat3286ce42018-09-10 21:33:08 +030070 iterable as it: object
71 Elements to divide into groups according to the key function.
72 key as keyfunc: object = None
73 A function for computing the group category for each element.
74 If the key function is not specified or is None, the element itself
75 is used for grouping.
76
77make an iterator that returns consecutive keys and groups from the iterable
78[clinic start generated code]*/
79
80static PyObject *
81itertools_groupby_impl(PyTypeObject *type, PyObject *it, PyObject *keyfunc)
82/*[clinic end generated code: output=cbb1ae3a90fd4141 input=6b3d123e87ff65a1]*/
83{
84 groupbyobject *gbo;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +000085
86 gbo = (groupbyobject *)type->tp_alloc(type, 0);
87 if (gbo == NULL)
88 return NULL;
89 gbo->tgtkey = NULL;
90 gbo->currkey = NULL;
91 gbo->currvalue = NULL;
92 gbo->keyfunc = keyfunc;
93 Py_INCREF(keyfunc);
94 gbo->it = PyObject_GetIter(it);
95 if (gbo->it == NULL) {
96 Py_DECREF(gbo);
97 return NULL;
98 }
99 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000100}
101
102static void
103groupby_dealloc(groupbyobject *gbo)
104{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000105 PyObject_GC_UnTrack(gbo);
106 Py_XDECREF(gbo->it);
107 Py_XDECREF(gbo->keyfunc);
108 Py_XDECREF(gbo->tgtkey);
109 Py_XDECREF(gbo->currkey);
110 Py_XDECREF(gbo->currvalue);
111 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000112}
113
114static int
115groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
116{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000117 Py_VISIT(gbo->it);
118 Py_VISIT(gbo->keyfunc);
119 Py_VISIT(gbo->tgtkey);
120 Py_VISIT(gbo->currkey);
121 Py_VISIT(gbo->currvalue);
122 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000123}
124
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300125Py_LOCAL_INLINE(int)
126groupby_step(groupbyobject *gbo)
127{
128 PyObject *newvalue, *newkey, *oldvalue;
129
130 newvalue = PyIter_Next(gbo->it);
131 if (newvalue == NULL)
132 return -1;
133
134 if (gbo->keyfunc == Py_None) {
135 newkey = newvalue;
136 Py_INCREF(newvalue);
137 } else {
Petr Viktorinffd97532020-02-11 17:46:57 +0100138 newkey = PyObject_CallOneArg(gbo->keyfunc, newvalue);
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300139 if (newkey == NULL) {
140 Py_DECREF(newvalue);
141 return -1;
142 }
143 }
144
145 oldvalue = gbo->currvalue;
146 gbo->currvalue = newvalue;
147 Py_XSETREF(gbo->currkey, newkey);
148 Py_XDECREF(oldvalue);
149 return 0;
150}
151
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000152static PyObject *
153groupby_next(groupbyobject *gbo)
154{
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300155 PyObject *r, *grouper;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000156
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300157 gbo->currgrouper = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000158 /* skip to next iteration group */
159 for (;;) {
160 if (gbo->currkey == NULL)
161 /* pass */;
162 else if (gbo->tgtkey == NULL)
163 break;
164 else {
165 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000166
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700167 rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000168 if (rcmp == -1)
169 return NULL;
170 else if (rcmp == 0)
171 break;
172 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000173
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300174 if (groupby_step(gbo) < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000175 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000176 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000177 Py_INCREF(gbo->currkey);
Serhiy Storchakaec397562016-04-06 09:50:03 +0300178 Py_XSETREF(gbo->tgtkey, gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000179
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000180 grouper = _grouper_create(gbo, gbo->tgtkey);
181 if (grouper == NULL)
182 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000183
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000184 r = PyTuple_Pack(2, gbo->currkey, grouper);
185 Py_DECREF(grouper);
186 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000187}
188
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000189static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530190groupby_reduce(groupbyobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000191{
192 /* reduce as a 'new' call with an optional 'setstate' if groupby
193 * has started
194 */
195 PyObject *value;
196 if (lz->tgtkey && lz->currkey && lz->currvalue)
197 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
198 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
199 else
200 value = Py_BuildValue("O(OO)", Py_TYPE(lz),
201 lz->it, lz->keyfunc);
202
203 return value;
204}
205
206PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
207
208static PyObject *
209groupby_setstate(groupbyobject *lz, PyObject *state)
210{
211 PyObject *currkey, *currvalue, *tgtkey;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300212 if (!PyTuple_Check(state)) {
213 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000214 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300215 }
216 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey)) {
217 return NULL;
218 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200219 Py_INCREF(currkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300220 Py_XSETREF(lz->currkey, currkey);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200221 Py_INCREF(currvalue);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300222 Py_XSETREF(lz->currvalue, currvalue);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200223 Py_INCREF(tgtkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300224 Py_XSETREF(lz->tgtkey, tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000225 Py_RETURN_NONE;
226}
227
228PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
229
230static PyMethodDef groupby_methods[] = {
231 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
232 reduce_doc},
233 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
234 setstate_doc},
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700235 {NULL, NULL} /* sentinel */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000236};
237
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000238static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000239 PyVarObject_HEAD_INIT(NULL, 0)
240 "itertools.groupby", /* tp_name */
241 sizeof(groupbyobject), /* tp_basicsize */
242 0, /* tp_itemsize */
243 /* methods */
244 (destructor)groupby_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200245 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000246 0, /* tp_getattr */
247 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200248 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000249 0, /* tp_repr */
250 0, /* tp_as_number */
251 0, /* tp_as_sequence */
252 0, /* tp_as_mapping */
253 0, /* tp_hash */
254 0, /* tp_call */
255 0, /* tp_str */
256 PyObject_GenericGetAttr, /* tp_getattro */
257 0, /* tp_setattro */
258 0, /* tp_as_buffer */
259 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
260 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einat3286ce42018-09-10 21:33:08 +0300261 itertools_groupby__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000262 (traverseproc)groupby_traverse, /* tp_traverse */
263 0, /* tp_clear */
264 0, /* tp_richcompare */
265 0, /* tp_weaklistoffset */
266 PyObject_SelfIter, /* tp_iter */
267 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000268 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000269 0, /* tp_members */
270 0, /* tp_getset */
271 0, /* tp_base */
272 0, /* tp_dict */
273 0, /* tp_descr_get */
274 0, /* tp_descr_set */
275 0, /* tp_dictoffset */
276 0, /* tp_init */
277 0, /* tp_alloc */
Tal Einat3286ce42018-09-10 21:33:08 +0300278 itertools_groupby, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000279 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000280};
281
282
283/* _grouper object (internal) ************************************************/
284
285typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000286 PyObject_HEAD
287 PyObject *parent;
288 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000289} _grouperobject;
290
Tal Einat3286ce42018-09-10 21:33:08 +0300291/*[clinic input]
292@classmethod
293itertools._grouper.__new__
294
295 parent: object(subclass_of='&groupby_type')
296 tgtkey: object
297 /
298[clinic start generated code]*/
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000299
300static PyObject *
Tal Einat3286ce42018-09-10 21:33:08 +0300301itertools__grouper_impl(PyTypeObject *type, PyObject *parent,
302 PyObject *tgtkey)
303/*[clinic end generated code: output=462efb1cdebb5914 input=dc180d7771fc8c59]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000304{
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000305 return _grouper_create((groupbyobject*) parent, tgtkey);
306}
307
308static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000309_grouper_create(groupbyobject *parent, PyObject *tgtkey)
310{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000311 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000312
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000313 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
314 if (igo == NULL)
315 return NULL;
316 igo->parent = (PyObject *)parent;
317 Py_INCREF(parent);
318 igo->tgtkey = tgtkey;
319 Py_INCREF(tgtkey);
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300320 parent->currgrouper = igo; /* borrowed reference */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000321
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000322 PyObject_GC_Track(igo);
323 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000324}
325
326static void
327_grouper_dealloc(_grouperobject *igo)
328{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000329 PyObject_GC_UnTrack(igo);
330 Py_DECREF(igo->parent);
331 Py_DECREF(igo->tgtkey);
332 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000333}
334
335static int
336_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
337{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000338 Py_VISIT(igo->parent);
339 Py_VISIT(igo->tgtkey);
340 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000341}
342
343static PyObject *
344_grouper_next(_grouperobject *igo)
345{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000346 groupbyobject *gbo = (groupbyobject *)igo->parent;
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300347 PyObject *r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000348 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000349
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300350 if (gbo->currgrouper != igo)
351 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000352 if (gbo->currvalue == NULL) {
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300353 if (groupby_step(gbo) < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000354 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000355 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000356
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000357 assert(gbo->currkey != NULL);
358 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
359 if (rcmp <= 0)
360 /* got any error or current group is end */
361 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000362
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000363 r = gbo->currvalue;
364 gbo->currvalue = NULL;
365 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000366
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000367 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000368}
369
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000370static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530371_grouper_reduce(_grouperobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000372{
Serhiy Storchakabb86bf42018-12-11 08:28:18 +0200373 _Py_IDENTIFIER(iter);
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300374 if (((groupbyobject *)lz->parent)->currgrouper != lz) {
Serhiy Storchakabb86bf42018-12-11 08:28:18 +0200375 return Py_BuildValue("N(())", _PyEval_GetBuiltinId(&PyId_iter));
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300376 }
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700377 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000378}
379
380static PyMethodDef _grouper_methods[] = {
381 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
382 reduce_doc},
383 {NULL, NULL} /* sentinel */
384};
385
386
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000387static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000388 PyVarObject_HEAD_INIT(NULL, 0)
389 "itertools._grouper", /* tp_name */
390 sizeof(_grouperobject), /* tp_basicsize */
391 0, /* tp_itemsize */
392 /* methods */
393 (destructor)_grouper_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200394 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000395 0, /* tp_getattr */
396 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200397 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000398 0, /* tp_repr */
399 0, /* tp_as_number */
400 0, /* tp_as_sequence */
401 0, /* tp_as_mapping */
402 0, /* tp_hash */
403 0, /* tp_call */
404 0, /* tp_str */
405 PyObject_GenericGetAttr, /* tp_getattro */
406 0, /* tp_setattro */
407 0, /* tp_as_buffer */
408 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
409 0, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700410 (traverseproc)_grouper_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000411 0, /* tp_clear */
412 0, /* tp_richcompare */
413 0, /* tp_weaklistoffset */
414 PyObject_SelfIter, /* tp_iter */
415 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000416 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000417 0, /* tp_members */
418 0, /* tp_getset */
419 0, /* tp_base */
420 0, /* tp_dict */
421 0, /* tp_descr_get */
422 0, /* tp_descr_set */
423 0, /* tp_dictoffset */
424 0, /* tp_init */
425 0, /* tp_alloc */
Tal Einat3286ce42018-09-10 21:33:08 +0300426 itertools__grouper, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000427 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000428};
429
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000430
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700431/* tee object and with supporting function and objects ***********************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000432
Raymond Hettingerad983e72003-11-12 14:32:26 +0000433/* The teedataobject pre-allocates space for LINKCELLS number of objects.
434 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000435 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000436 the other structure members including PyHEAD overhead. The larger the
437 value, the less memory overhead per object and the less time spent
438 allocating/deallocating new links. The smaller the number, the less
439 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000440*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000441#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000442
443typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000444 PyObject_HEAD
445 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700446 int numread; /* 0 <= numread <= LINKCELLS */
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300447 int running;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000448 PyObject *nextlink;
449 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000450} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000451
452typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000453 PyObject_HEAD
454 teedataobject *dataobj;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700455 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000456 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000457} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000458
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000459static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000460teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000461{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000462 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000463
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000464 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
465 if (tdo == NULL)
466 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000467
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300468 tdo->running = 0;
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);
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300497 if (tdo->running) {
498 PyErr_SetString(PyExc_RuntimeError,
499 "cannot re-enter the tee iterator");
500 return NULL;
501 }
502 tdo->running = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000503 value = PyIter_Next(tdo->it);
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300504 tdo->running = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000505 if (value == NULL)
506 return NULL;
507 tdo->numread++;
508 tdo->values[i] = value;
509 }
510 Py_INCREF(value);
511 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000512}
513
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000514static int
515teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
516{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000517 int i;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700518
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000519 Py_VISIT(tdo->it);
520 for (i = 0; i < tdo->numread; i++)
521 Py_VISIT(tdo->values[i]);
522 Py_VISIT(tdo->nextlink);
523 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000524}
525
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200526static void
527teedataobject_safe_decref(PyObject *obj)
528{
Andy Lesterdffe4c02020-03-04 07:15:20 -0600529 while (obj && Py_IS_TYPE(obj, &teedataobject_type) &&
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200530 Py_REFCNT(obj) == 1) {
531 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
532 ((teedataobject *)obj)->nextlink = NULL;
533 Py_DECREF(obj);
534 obj = nextlink;
535 }
536 Py_XDECREF(obj);
537}
538
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000539static int
540teedataobject_clear(teedataobject *tdo)
541{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000542 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200543 PyObject *tmp;
544
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000545 Py_CLEAR(tdo->it);
546 for (i=0 ; i<tdo->numread ; i++)
547 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200548 tmp = tdo->nextlink;
549 tdo->nextlink = NULL;
550 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000551 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000552}
553
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000554static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000555teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000556{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000557 PyObject_GC_UnTrack(tdo);
558 teedataobject_clear(tdo);
559 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000560}
561
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000562static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530563teedataobject_reduce(teedataobject *tdo, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000564{
565 int i;
566 /* create a temporary list of already iterated values */
567 PyObject *values = PyList_New(tdo->numread);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700568
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000569 if (!values)
570 return NULL;
571 for (i=0 ; i<tdo->numread ; i++) {
572 Py_INCREF(tdo->values[i]);
573 PyList_SET_ITEM(values, i, tdo->values[i]);
574 }
575 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
576 values,
577 tdo->nextlink ? tdo->nextlink : Py_None);
578}
579
Tal Einatc4bccd32018-09-12 00:49:13 +0300580/*[clinic input]
581@classmethod
582itertools.teedataobject.__new__
583 iterable as it: object
584 values: object(subclass_of='&PyList_Type')
585 next: object
586 /
587Data container common to multiple tee objects.
588[clinic start generated code]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000589
590static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300591itertools_teedataobject_impl(PyTypeObject *type, PyObject *it,
592 PyObject *values, PyObject *next)
593/*[clinic end generated code: output=3343ceb07e08df5e input=be60f2fabd2b72ba]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000594{
595 teedataobject *tdo;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000596 Py_ssize_t i, len;
597
598 assert(type == &teedataobject_type);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000599
600 tdo = (teedataobject *)teedataobject_newinternal(it);
601 if (!tdo)
602 return NULL;
603
604 len = PyList_GET_SIZE(values);
605 if (len > LINKCELLS)
606 goto err;
607 for (i=0; i<len; i++) {
608 tdo->values[i] = PyList_GET_ITEM(values, i);
609 Py_INCREF(tdo->values[i]);
610 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200611 /* len <= LINKCELLS < INT_MAX */
612 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000613
614 if (len == LINKCELLS) {
615 if (next != Py_None) {
Andy Lester55728702020-03-06 16:53:17 -0600616 if (!Py_IS_TYPE(next, &teedataobject_type))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000617 goto err;
618 assert(tdo->nextlink == NULL);
619 Py_INCREF(next);
620 tdo->nextlink = next;
621 }
622 } else {
623 if (next != Py_None)
624 goto err; /* shouldn't have a next if we are not full */
625 }
626 return (PyObject*)tdo;
627
628err:
629 Py_XDECREF(tdo);
630 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
631 return NULL;
632}
633
634static PyMethodDef teedataobject_methods[] = {
635 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
636 reduce_doc},
637 {NULL, NULL} /* sentinel */
638};
639
Raymond Hettingerad983e72003-11-12 14:32:26 +0000640static PyTypeObject teedataobject_type = {
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700641 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000642 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000643 sizeof(teedataobject), /* tp_basicsize */
644 0, /* tp_itemsize */
645 /* methods */
646 (destructor)teedataobject_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200647 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000648 0, /* tp_getattr */
649 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200650 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000651 0, /* tp_repr */
652 0, /* tp_as_number */
653 0, /* tp_as_sequence */
654 0, /* tp_as_mapping */
655 0, /* tp_hash */
656 0, /* tp_call */
657 0, /* tp_str */
658 PyObject_GenericGetAttr, /* tp_getattro */
659 0, /* tp_setattro */
660 0, /* tp_as_buffer */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700661 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300662 itertools_teedataobject__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000663 (traverseproc)teedataobject_traverse, /* tp_traverse */
664 (inquiry)teedataobject_clear, /* tp_clear */
665 0, /* tp_richcompare */
666 0, /* tp_weaklistoffset */
667 0, /* tp_iter */
668 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000669 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000670 0, /* tp_members */
671 0, /* tp_getset */
672 0, /* tp_base */
673 0, /* tp_dict */
674 0, /* tp_descr_get */
675 0, /* tp_descr_set */
676 0, /* tp_dictoffset */
677 0, /* tp_init */
678 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +0300679 itertools_teedataobject, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000680 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000681};
682
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000683
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000684static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000685tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000686{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000687 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000688
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;
722 PyObject_GC_Track(newto);
723 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000724}
725
726PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
727
728static PyObject *
729tee_fromiterable(PyObject *iterable)
730{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000731 teeobject *to;
732 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000733
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000734 it = PyObject_GetIter(iterable);
735 if (it == NULL)
736 return NULL;
737 if (PyObject_TypeCheck(it, &tee_type)) {
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530738 to = (teeobject *)tee_copy((teeobject *)it, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000739 goto done;
740 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000741
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000742 to = PyObject_GC_New(teeobject, &tee_type);
743 if (to == NULL)
744 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000745 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000746 if (!to->dataobj) {
747 PyObject_GC_Del(to);
748 to = NULL;
749 goto done;
750 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000751
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000752 to->index = 0;
753 to->weakreflist = NULL;
754 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000755done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000756 Py_XDECREF(it);
757 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000758}
759
Tal Einatc4bccd32018-09-12 00:49:13 +0300760/*[clinic input]
761@classmethod
762itertools._tee.__new__
763 iterable: object
764 /
765Iterator wrapped to make it copyable.
766[clinic start generated code]*/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000767
Tal Einatc4bccd32018-09-12 00:49:13 +0300768static PyObject *
769itertools__tee_impl(PyTypeObject *type, PyObject *iterable)
770/*[clinic end generated code: output=b02d3fd26c810c3f input=adc0779d2afe37a2]*/
771{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000772 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000773}
774
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000775static int
776tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000777{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000778 if (to->weakreflist != NULL)
779 PyObject_ClearWeakRefs((PyObject *) to);
780 Py_CLEAR(to->dataobj);
781 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000782}
783
784static void
785tee_dealloc(teeobject *to)
786{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000787 PyObject_GC_UnTrack(to);
788 tee_clear(to);
789 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000790}
791
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000792static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530793tee_reduce(teeobject *to, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000794{
795 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
796}
797
798static PyObject *
799tee_setstate(teeobject *to, PyObject *state)
800{
801 teedataobject *tdo;
802 int index;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300803 if (!PyTuple_Check(state)) {
804 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000805 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300806 }
807 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {
808 return NULL;
809 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000810 if (index < 0 || index > LINKCELLS) {
811 PyErr_SetString(PyExc_ValueError, "Index out of range");
812 return NULL;
813 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200814 Py_INCREF(tdo);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300815 Py_XSETREF(to->dataobj, tdo);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000816 to->index = index;
817 Py_RETURN_NONE;
818}
819
Raymond Hettingerad983e72003-11-12 14:32:26 +0000820static PyMethodDef tee_methods[] = {
Raymond Hettingera6a2d442015-08-16 14:38:07 -0700821 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
822 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
823 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000824 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000825};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000826
827static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000828 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000829 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000830 sizeof(teeobject), /* tp_basicsize */
831 0, /* tp_itemsize */
832 /* methods */
833 (destructor)tee_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200834 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000835 0, /* tp_getattr */
836 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200837 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000838 0, /* tp_repr */
839 0, /* tp_as_number */
840 0, /* tp_as_sequence */
841 0, /* tp_as_mapping */
842 0, /* tp_hash */
843 0, /* tp_call */
844 0, /* tp_str */
845 0, /* tp_getattro */
846 0, /* tp_setattro */
847 0, /* tp_as_buffer */
848 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300849 itertools__tee__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000850 (traverseproc)tee_traverse, /* tp_traverse */
851 (inquiry)tee_clear, /* tp_clear */
852 0, /* tp_richcompare */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700853 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000854 PyObject_SelfIter, /* tp_iter */
855 (iternextfunc)tee_next, /* tp_iternext */
856 tee_methods, /* tp_methods */
857 0, /* tp_members */
858 0, /* tp_getset */
859 0, /* tp_base */
860 0, /* tp_dict */
861 0, /* tp_descr_get */
862 0, /* tp_descr_set */
863 0, /* tp_dictoffset */
864 0, /* tp_init */
865 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +0300866 itertools__tee, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000867 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000868};
869
Tal Einatc4bccd32018-09-12 00:49:13 +0300870/*[clinic input]
871itertools.tee
872 iterable: object
873 n: Py_ssize_t = 2
874 /
875Returns a tuple of n independent iterators.
876[clinic start generated code]*/
877
Raymond Hettingerad983e72003-11-12 14:32:26 +0000878static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300879itertools_tee_impl(PyObject *module, PyObject *iterable, Py_ssize_t n)
880/*[clinic end generated code: output=1c64519cd859c2f0 input=c99a1472c425d66d]*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000881{
Tal Einatc4bccd32018-09-12 00:49:13 +0300882 Py_ssize_t i;
883 PyObject *it, *copyable, *copyfunc, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +0200884 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000885
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000886 if (n < 0) {
887 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
888 return NULL;
889 }
890 result = PyTuple_New(n);
891 if (result == NULL)
892 return NULL;
893 if (n == 0)
894 return result;
895 it = PyObject_GetIter(iterable);
896 if (it == NULL) {
897 Py_DECREF(result);
898 return NULL;
899 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200900
Serhiy Storchakaf320be72018-01-25 10:49:40 +0200901 if (_PyObject_LookupAttrId(it, &PyId___copy__, &copyfunc) < 0) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200902 Py_DECREF(it);
903 Py_DECREF(result);
904 return NULL;
905 }
Serhiy Storchakaf320be72018-01-25 10:49:40 +0200906 if (copyfunc != NULL) {
907 copyable = it;
908 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200909 else {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000910 copyable = tee_fromiterable(it);
911 Py_DECREF(it);
912 if (copyable == NULL) {
913 Py_DECREF(result);
914 return NULL;
915 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200916 copyfunc = _PyObject_GetAttrId(copyable, &PyId___copy__);
917 if (copyfunc == NULL) {
918 Py_DECREF(copyable);
919 Py_DECREF(result);
920 return NULL;
921 }
922 }
Martin v. Löwisafe55bb2011-10-09 10:38:36 +0200923
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200924 PyTuple_SET_ITEM(result, 0, copyable);
925 for (i = 1; i < n; i++) {
926 copyable = _PyObject_CallNoArg(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000927 if (copyable == NULL) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200928 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000929 Py_DECREF(result);
930 return NULL;
931 }
932 PyTuple_SET_ITEM(result, i, copyable);
933 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +0200934 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000935 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000936}
937
Raymond Hettingerad983e72003-11-12 14:32:26 +0000938
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700939/* cycle object **************************************************************/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000940
941typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000942 PyObject_HEAD
943 PyObject *it;
944 PyObject *saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700945 Py_ssize_t index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000946 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000947} cycleobject;
948
Tal Einatc4bccd32018-09-12 00:49:13 +0300949/*[clinic input]
950@classmethod
951itertools.cycle.__new__
952 iterable: object
953 /
954Return elements from the iterable until it is exhausted. Then repeat the sequence indefinitely.
955[clinic start generated code]*/
956
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000957static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300958itertools_cycle_impl(PyTypeObject *type, PyObject *iterable)
959/*[clinic end generated code: output=f60e5ec17a45b35c input=9d1d84bcf66e908b]*/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000960{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000961 PyObject *it;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000962 PyObject *saved;
963 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000964
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000965 /* Get iterator. */
966 it = PyObject_GetIter(iterable);
967 if (it == NULL)
968 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000969
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000970 saved = PyList_New(0);
971 if (saved == NULL) {
972 Py_DECREF(it);
973 return NULL;
974 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000975
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000976 /* create cycleobject structure */
977 lz = (cycleobject *)type->tp_alloc(type, 0);
978 if (lz == NULL) {
979 Py_DECREF(it);
980 Py_DECREF(saved);
981 return NULL;
982 }
983 lz->it = it;
984 lz->saved = saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700985 lz->index = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000986 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000987
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000988 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000989}
990
991static void
992cycle_dealloc(cycleobject *lz)
993{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000994 PyObject_GC_UnTrack(lz);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000995 Py_XDECREF(lz->it);
Raymond Hettingerca3788c2015-08-16 14:49:24 -0700996 Py_XDECREF(lz->saved);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000997 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000998}
999
1000static int
1001cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
1002{
Hai Shi47a23fc2020-06-07 20:05:36 +08001003 Py_VISIT(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001004 Py_VISIT(lz->saved);
1005 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001006}
1007
1008static PyObject *
1009cycle_next(cycleobject *lz)
1010{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001011 PyObject *item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001012
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001013 if (lz->it != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001014 item = PyIter_Next(lz->it);
1015 if (item != NULL) {
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001016 if (lz->firstpass)
1017 return item;
1018 if (PyList_Append(lz->saved, item)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001019 Py_DECREF(item);
1020 return NULL;
1021 }
1022 return item;
1023 }
Raymond Hettinger98958fe2015-08-15 15:09:30 -07001024 /* Note: StopIteration is already cleared by PyIter_Next() */
1025 if (PyErr_Occurred())
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001026 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001027 Py_CLEAR(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001028 }
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001029 if (PyList_GET_SIZE(lz->saved) == 0)
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001030 return NULL;
1031 item = PyList_GET_ITEM(lz->saved, lz->index);
1032 lz->index++;
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001033 if (lz->index >= PyList_GET_SIZE(lz->saved))
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001034 lz->index = 0;
1035 Py_INCREF(item);
1036 return item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001037}
1038
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001039static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301040cycle_reduce(cycleobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001041{
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001042 /* Create a new cycle with the iterator tuple, then set the saved state */
1043 if (lz->it == NULL) {
1044 PyObject *it = PyObject_GetIter(lz->saved);
1045 if (it == NULL)
1046 return NULL;
1047 if (lz->index != 0) {
1048 _Py_IDENTIFIER(__setstate__);
1049 PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
1050 "n", lz->index);
1051 if (res == NULL) {
1052 Py_DECREF(it);
1053 return NULL;
1054 }
1055 Py_DECREF(res);
1056 }
Serhiy Storchaka1f21eaa2019-09-01 12:16:51 +03001057 return Py_BuildValue("O(N)(OO)", Py_TYPE(lz), it, lz->saved, Py_True);
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001058 }
Serhiy Storchaka1f21eaa2019-09-01 12:16:51 +03001059 return Py_BuildValue("O(O)(OO)", Py_TYPE(lz), lz->it, lz->saved,
1060 lz->firstpass ? Py_True : Py_False);
Raymond Hettingerc39786d2015-08-15 15:16:12 -07001061}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001062
1063static PyObject *
1064cycle_setstate(cycleobject *lz, PyObject *state)
1065{
1066 PyObject *saved=NULL;
1067 int firstpass;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001068 if (!PyTuple_Check(state)) {
1069 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001070 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001071 }
1072 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
1073 return NULL;
1074 }
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001075 Py_INCREF(saved);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001076 Py_XSETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001077 lz->firstpass = firstpass != 0;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001078 lz->index = 0;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001079 Py_RETURN_NONE;
1080}
1081
1082static PyMethodDef cycle_methods[] = {
1083 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
1084 reduce_doc},
1085 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
1086 setstate_doc},
1087 {NULL, NULL} /* sentinel */
1088};
1089
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001090static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001091 PyVarObject_HEAD_INIT(NULL, 0)
1092 "itertools.cycle", /* tp_name */
1093 sizeof(cycleobject), /* tp_basicsize */
1094 0, /* tp_itemsize */
1095 /* methods */
1096 (destructor)cycle_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001097 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001098 0, /* tp_getattr */
1099 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001100 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001101 0, /* tp_repr */
1102 0, /* tp_as_number */
1103 0, /* tp_as_sequence */
1104 0, /* tp_as_mapping */
1105 0, /* tp_hash */
1106 0, /* tp_call */
1107 0, /* tp_str */
1108 PyObject_GenericGetAttr, /* tp_getattro */
1109 0, /* tp_setattro */
1110 0, /* tp_as_buffer */
1111 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1112 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001113 itertools_cycle__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001114 (traverseproc)cycle_traverse, /* tp_traverse */
1115 0, /* tp_clear */
1116 0, /* tp_richcompare */
1117 0, /* tp_weaklistoffset */
1118 PyObject_SelfIter, /* tp_iter */
1119 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001120 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001121 0, /* tp_members */
1122 0, /* tp_getset */
1123 0, /* tp_base */
1124 0, /* tp_dict */
1125 0, /* tp_descr_get */
1126 0, /* tp_descr_set */
1127 0, /* tp_dictoffset */
1128 0, /* tp_init */
1129 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001130 itertools_cycle, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001131 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001132};
1133
1134
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001135/* dropwhile object **********************************************************/
1136
1137typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001138 PyObject_HEAD
1139 PyObject *func;
1140 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001141 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001142} dropwhileobject;
1143
Tal Einatc4bccd32018-09-12 00:49:13 +03001144/*[clinic input]
1145@classmethod
1146itertools.dropwhile.__new__
1147 predicate as func: object
1148 iterable as seq: object
1149 /
1150Drop items from the iterable while predicate(item) is true.
1151
1152Afterwards, return every element until the iterable is exhausted.
1153[clinic start generated code]*/
1154
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001155static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001156itertools_dropwhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1157/*[clinic end generated code: output=92f9d0d89af149e4 input=d39737147c9f0a26]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001158{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001159 PyObject *it;
1160 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001161
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001162 /* Get iterator. */
1163 it = PyObject_GetIter(seq);
1164 if (it == NULL)
1165 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001166
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001167 /* create dropwhileobject structure */
1168 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1169 if (lz == NULL) {
1170 Py_DECREF(it);
1171 return NULL;
1172 }
1173 Py_INCREF(func);
1174 lz->func = func;
1175 lz->it = it;
1176 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001177
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001178 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001179}
1180
1181static void
1182dropwhile_dealloc(dropwhileobject *lz)
1183{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001184 PyObject_GC_UnTrack(lz);
1185 Py_XDECREF(lz->func);
1186 Py_XDECREF(lz->it);
1187 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001188}
1189
1190static int
1191dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1192{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001193 Py_VISIT(lz->it);
1194 Py_VISIT(lz->func);
1195 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001196}
1197
1198static PyObject *
1199dropwhile_next(dropwhileobject *lz)
1200{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001201 PyObject *item, *good;
1202 PyObject *it = lz->it;
1203 long ok;
1204 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001205
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001206 iternext = *Py_TYPE(it)->tp_iternext;
1207 for (;;) {
1208 item = iternext(it);
1209 if (item == NULL)
1210 return NULL;
1211 if (lz->start == 1)
1212 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001213
Petr Viktorinffd97532020-02-11 17:46:57 +01001214 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001215 if (good == NULL) {
1216 Py_DECREF(item);
1217 return NULL;
1218 }
1219 ok = PyObject_IsTrue(good);
1220 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001221 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001222 lz->start = 1;
1223 return item;
1224 }
1225 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001226 if (ok < 0)
1227 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001228 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001229}
1230
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001231static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301232dropwhile_reduce(dropwhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001233{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001234 return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->start);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001235}
1236
1237static PyObject *
1238dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1239{
1240 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001241 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001242 return NULL;
1243 lz->start = start;
1244 Py_RETURN_NONE;
1245}
1246
1247static PyMethodDef dropwhile_methods[] = {
1248 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1249 reduce_doc},
1250 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1251 setstate_doc},
1252 {NULL, NULL} /* sentinel */
1253};
1254
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001255static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001256 PyVarObject_HEAD_INIT(NULL, 0)
1257 "itertools.dropwhile", /* tp_name */
1258 sizeof(dropwhileobject), /* tp_basicsize */
1259 0, /* tp_itemsize */
1260 /* methods */
1261 (destructor)dropwhile_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001262 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001263 0, /* tp_getattr */
1264 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001265 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001266 0, /* tp_repr */
1267 0, /* tp_as_number */
1268 0, /* tp_as_sequence */
1269 0, /* tp_as_mapping */
1270 0, /* tp_hash */
1271 0, /* tp_call */
1272 0, /* tp_str */
1273 PyObject_GenericGetAttr, /* tp_getattro */
1274 0, /* tp_setattro */
1275 0, /* tp_as_buffer */
1276 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1277 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001278 itertools_dropwhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001279 (traverseproc)dropwhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001280 0, /* tp_clear */
1281 0, /* tp_richcompare */
1282 0, /* tp_weaklistoffset */
1283 PyObject_SelfIter, /* tp_iter */
1284 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001285 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001286 0, /* tp_members */
1287 0, /* tp_getset */
1288 0, /* tp_base */
1289 0, /* tp_dict */
1290 0, /* tp_descr_get */
1291 0, /* tp_descr_set */
1292 0, /* tp_dictoffset */
1293 0, /* tp_init */
1294 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001295 itertools_dropwhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001296 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001297};
1298
1299
1300/* takewhile object **********************************************************/
1301
1302typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001303 PyObject_HEAD
1304 PyObject *func;
1305 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001306 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001307} takewhileobject;
1308
Tal Einatc4bccd32018-09-12 00:49:13 +03001309/*[clinic input]
1310@classmethod
1311itertools.takewhile.__new__
1312 predicate as func: object
1313 iterable as seq: object
1314 /
1315Return successive entries from an iterable as long as the predicate evaluates to true for each entry.
1316[clinic start generated code]*/
1317
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001318static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001319itertools_takewhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1320/*[clinic end generated code: output=bb179ea7864e2ef6 input=ba5255f7519aa119]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001321{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001322 PyObject *it;
1323 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001324
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001325 /* Get iterator. */
1326 it = PyObject_GetIter(seq);
1327 if (it == NULL)
1328 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001329
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001330 /* create takewhileobject structure */
1331 lz = (takewhileobject *)type->tp_alloc(type, 0);
1332 if (lz == NULL) {
1333 Py_DECREF(it);
1334 return NULL;
1335 }
1336 Py_INCREF(func);
1337 lz->func = func;
1338 lz->it = it;
1339 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001340
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001341 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001342}
1343
1344static void
1345takewhile_dealloc(takewhileobject *lz)
1346{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001347 PyObject_GC_UnTrack(lz);
1348 Py_XDECREF(lz->func);
1349 Py_XDECREF(lz->it);
1350 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001351}
1352
1353static int
1354takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1355{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001356 Py_VISIT(lz->it);
1357 Py_VISIT(lz->func);
1358 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001359}
1360
1361static PyObject *
1362takewhile_next(takewhileobject *lz)
1363{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001364 PyObject *item, *good;
1365 PyObject *it = lz->it;
1366 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001367
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001368 if (lz->stop == 1)
1369 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001370
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001371 item = (*Py_TYPE(it)->tp_iternext)(it);
1372 if (item == NULL)
1373 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001374
Petr Viktorinffd97532020-02-11 17:46:57 +01001375 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001376 if (good == NULL) {
1377 Py_DECREF(item);
1378 return NULL;
1379 }
1380 ok = PyObject_IsTrue(good);
1381 Py_DECREF(good);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001382 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001383 return item;
1384 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001385 if (ok == 0)
1386 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001387 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001388}
1389
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001390static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301391takewhile_reduce(takewhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001392{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001393 return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->stop);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001394}
1395
1396static PyObject *
1397takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1398{
1399 int stop = PyObject_IsTrue(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001400
Antoine Pitrou721738f2012-08-15 23:20:39 +02001401 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001402 return NULL;
1403 lz->stop = stop;
1404 Py_RETURN_NONE;
1405}
1406
1407static PyMethodDef takewhile_reduce_methods[] = {
1408 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1409 reduce_doc},
1410 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1411 setstate_doc},
1412 {NULL, NULL} /* sentinel */
1413};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001414
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001415static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001416 PyVarObject_HEAD_INIT(NULL, 0)
1417 "itertools.takewhile", /* tp_name */
1418 sizeof(takewhileobject), /* tp_basicsize */
1419 0, /* tp_itemsize */
1420 /* methods */
1421 (destructor)takewhile_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001422 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001423 0, /* tp_getattr */
1424 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001425 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001426 0, /* tp_repr */
1427 0, /* tp_as_number */
1428 0, /* tp_as_sequence */
1429 0, /* tp_as_mapping */
1430 0, /* tp_hash */
1431 0, /* tp_call */
1432 0, /* tp_str */
1433 PyObject_GenericGetAttr, /* tp_getattro */
1434 0, /* tp_setattro */
1435 0, /* tp_as_buffer */
1436 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1437 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001438 itertools_takewhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001439 (traverseproc)takewhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001440 0, /* tp_clear */
1441 0, /* tp_richcompare */
1442 0, /* tp_weaklistoffset */
1443 PyObject_SelfIter, /* tp_iter */
1444 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001445 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001446 0, /* tp_members */
1447 0, /* tp_getset */
1448 0, /* tp_base */
1449 0, /* tp_dict */
1450 0, /* tp_descr_get */
1451 0, /* tp_descr_set */
1452 0, /* tp_dictoffset */
1453 0, /* tp_init */
1454 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001455 itertools_takewhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001456 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457};
1458
1459
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001460/* islice object *************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001461
1462typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001463 PyObject_HEAD
1464 PyObject *it;
1465 Py_ssize_t next;
1466 Py_ssize_t stop;
1467 Py_ssize_t step;
1468 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001469} isliceobject;
1470
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001471static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001472
1473static PyObject *
1474islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1475{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001476 PyObject *seq;
1477 Py_ssize_t start=0, stop=-1, step=1;
1478 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1479 Py_ssize_t numargs;
1480 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001481
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001482 if (type == &islice_type && !_PyArg_NoKeywords("islice", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001483 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001484
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001485 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1486 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001487
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001488 numargs = PyTuple_Size(args);
1489 if (numargs == 2) {
1490 if (a1 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001491 stop = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001492 if (stop == -1) {
1493 if (PyErr_Occurred())
1494 PyErr_Clear();
1495 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001496 "Stop argument for islice() must be None or "
1497 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001498 return NULL;
1499 }
1500 }
1501 } else {
1502 if (a1 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001503 start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001504 if (start == -1 && PyErr_Occurred())
1505 PyErr_Clear();
1506 if (a2 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001507 stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001508 if (stop == -1) {
1509 if (PyErr_Occurred())
1510 PyErr_Clear();
1511 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001512 "Stop argument for islice() must be None or "
1513 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001514 return NULL;
1515 }
1516 }
1517 }
1518 if (start<0 || stop<-1) {
1519 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001520 "Indices for islice() must be None or "
1521 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001522 return NULL;
1523 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001524
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001525 if (a3 != NULL) {
1526 if (a3 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001527 step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001528 if (step == -1 && PyErr_Occurred())
1529 PyErr_Clear();
1530 }
1531 if (step<1) {
1532 PyErr_SetString(PyExc_ValueError,
1533 "Step for islice() must be a positive integer or None.");
1534 return NULL;
1535 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001536
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001537 /* Get iterator. */
1538 it = PyObject_GetIter(seq);
1539 if (it == NULL)
1540 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001541
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001542 /* create isliceobject structure */
1543 lz = (isliceobject *)type->tp_alloc(type, 0);
1544 if (lz == NULL) {
1545 Py_DECREF(it);
1546 return NULL;
1547 }
1548 lz->it = it;
1549 lz->next = start;
1550 lz->stop = stop;
1551 lz->step = step;
1552 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001553
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001554 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001555}
1556
1557static void
1558islice_dealloc(isliceobject *lz)
1559{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001560 PyObject_GC_UnTrack(lz);
1561 Py_XDECREF(lz->it);
1562 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001563}
1564
1565static int
1566islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1567{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001568 Py_VISIT(lz->it);
1569 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001570}
1571
1572static PyObject *
1573islice_next(isliceobject *lz)
1574{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001575 PyObject *item;
1576 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001577 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001578 Py_ssize_t oldnext;
1579 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001580
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001581 if (it == NULL)
1582 return NULL;
1583
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001584 iternext = *Py_TYPE(it)->tp_iternext;
1585 while (lz->cnt < lz->next) {
1586 item = iternext(it);
1587 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001588 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001589 Py_DECREF(item);
1590 lz->cnt++;
1591 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001592 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001593 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001594 item = iternext(it);
1595 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001596 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001597 lz->cnt++;
1598 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001599 /* The (size_t) cast below avoids the danger of undefined
1600 behaviour from signed integer overflow. */
1601 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001602 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1603 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001604 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001605
1606empty:
1607 Py_CLEAR(lz->it);
1608 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001609}
1610
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001611static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301612islice_reduce(isliceobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001613{
1614 /* When unpickled, generate a new object with the same bounds,
1615 * then 'setstate' with the next and count
1616 */
1617 PyObject *stop;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001618
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001619 if (lz->it == NULL) {
1620 PyObject *empty_list;
1621 PyObject *empty_it;
1622 empty_list = PyList_New(0);
1623 if (empty_list == NULL)
1624 return NULL;
1625 empty_it = PyObject_GetIter(empty_list);
1626 Py_DECREF(empty_list);
1627 if (empty_it == NULL)
1628 return NULL;
1629 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1630 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001631 if (lz->stop == -1) {
1632 stop = Py_None;
1633 Py_INCREF(stop);
1634 } else {
1635 stop = PyLong_FromSsize_t(lz->stop);
1636 if (stop == NULL)
1637 return NULL;
1638 }
1639 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1640 lz->it, lz->next, stop, lz->step,
1641 lz->cnt);
1642}
1643
1644static PyObject *
1645islice_setstate(isliceobject *lz, PyObject *state)
1646{
1647 Py_ssize_t cnt = PyLong_AsSsize_t(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001648
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001649 if (cnt == -1 && PyErr_Occurred())
1650 return NULL;
1651 lz->cnt = cnt;
1652 Py_RETURN_NONE;
1653}
1654
1655static PyMethodDef islice_methods[] = {
1656 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1657 reduce_doc},
1658 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1659 setstate_doc},
1660 {NULL, NULL} /* sentinel */
1661};
1662
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001663PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001664"islice(iterable, stop) --> islice object\n\
1665islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001666\n\
1667Return an iterator whose next() method returns selected values from an\n\
1668iterable. If start is specified, will skip all preceding elements;\n\
1669otherwise, start defaults to zero. Step defaults to one. If\n\
oldkaa0735f2018-02-02 16:52:55 +08001670specified as another value, step determines how many values are\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001671skipped between successive calls. Works like a slice() on a list\n\
1672but returns an iterator.");
1673
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001674static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001675 PyVarObject_HEAD_INIT(NULL, 0)
1676 "itertools.islice", /* tp_name */
1677 sizeof(isliceobject), /* tp_basicsize */
1678 0, /* tp_itemsize */
1679 /* methods */
1680 (destructor)islice_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001681 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001682 0, /* tp_getattr */
1683 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001684 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001685 0, /* tp_repr */
1686 0, /* tp_as_number */
1687 0, /* tp_as_sequence */
1688 0, /* tp_as_mapping */
1689 0, /* tp_hash */
1690 0, /* tp_call */
1691 0, /* tp_str */
1692 PyObject_GenericGetAttr, /* tp_getattro */
1693 0, /* tp_setattro */
1694 0, /* tp_as_buffer */
1695 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1696 Py_TPFLAGS_BASETYPE, /* tp_flags */
1697 islice_doc, /* tp_doc */
1698 (traverseproc)islice_traverse, /* tp_traverse */
1699 0, /* tp_clear */
1700 0, /* tp_richcompare */
1701 0, /* tp_weaklistoffset */
1702 PyObject_SelfIter, /* tp_iter */
1703 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001704 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001705 0, /* tp_members */
1706 0, /* tp_getset */
1707 0, /* tp_base */
1708 0, /* tp_dict */
1709 0, /* tp_descr_get */
1710 0, /* tp_descr_set */
1711 0, /* tp_dictoffset */
1712 0, /* tp_init */
1713 0, /* tp_alloc */
1714 islice_new, /* tp_new */
1715 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001716};
1717
1718
1719/* starmap object ************************************************************/
1720
1721typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001722 PyObject_HEAD
1723 PyObject *func;
1724 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001725} starmapobject;
1726
Tal Einatc4bccd32018-09-12 00:49:13 +03001727/*[clinic input]
1728@classmethod
1729itertools.starmap.__new__
1730 function as func: object
1731 iterable as seq: object
1732 /
1733Return an iterator whose values are returned from the function evaluated with an argument tuple taken from the given sequence.
1734[clinic start generated code]*/
1735
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001736static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001737itertools_starmap_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1738/*[clinic end generated code: output=79eeb81d452c6e8d input=844766df6a0d4dad]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001739{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001740 PyObject *it;
1741 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001742
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001743 /* Get iterator. */
1744 it = PyObject_GetIter(seq);
1745 if (it == NULL)
1746 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001747
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001748 /* create starmapobject structure */
1749 lz = (starmapobject *)type->tp_alloc(type, 0);
1750 if (lz == NULL) {
1751 Py_DECREF(it);
1752 return NULL;
1753 }
1754 Py_INCREF(func);
1755 lz->func = func;
1756 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001757
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001758 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001759}
1760
1761static void
1762starmap_dealloc(starmapobject *lz)
1763{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001764 PyObject_GC_UnTrack(lz);
1765 Py_XDECREF(lz->func);
1766 Py_XDECREF(lz->it);
1767 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001768}
1769
1770static int
1771starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1772{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001773 Py_VISIT(lz->it);
1774 Py_VISIT(lz->func);
1775 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001776}
1777
1778static PyObject *
1779starmap_next(starmapobject *lz)
1780{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001781 PyObject *args;
1782 PyObject *result;
1783 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001784
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001785 args = (*Py_TYPE(it)->tp_iternext)(it);
1786 if (args == NULL)
1787 return NULL;
1788 if (!PyTuple_CheckExact(args)) {
1789 PyObject *newargs = PySequence_Tuple(args);
1790 Py_DECREF(args);
1791 if (newargs == NULL)
1792 return NULL;
1793 args = newargs;
1794 }
1795 result = PyObject_Call(lz->func, args, NULL);
1796 Py_DECREF(args);
1797 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001798}
1799
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001800static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301801starmap_reduce(starmapobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001802{
1803 /* Just pickle the iterator */
1804 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1805}
1806
1807static PyMethodDef starmap_methods[] = {
1808 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1809 reduce_doc},
1810 {NULL, NULL} /* sentinel */
1811};
1812
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001813static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001814 PyVarObject_HEAD_INIT(NULL, 0)
1815 "itertools.starmap", /* tp_name */
1816 sizeof(starmapobject), /* tp_basicsize */
1817 0, /* tp_itemsize */
1818 /* methods */
1819 (destructor)starmap_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001820 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001821 0, /* tp_getattr */
1822 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001823 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001824 0, /* tp_repr */
1825 0, /* tp_as_number */
1826 0, /* tp_as_sequence */
1827 0, /* tp_as_mapping */
1828 0, /* tp_hash */
1829 0, /* tp_call */
1830 0, /* tp_str */
1831 PyObject_GenericGetAttr, /* tp_getattro */
1832 0, /* tp_setattro */
1833 0, /* tp_as_buffer */
1834 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1835 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001836 itertools_starmap__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001837 (traverseproc)starmap_traverse, /* tp_traverse */
1838 0, /* tp_clear */
1839 0, /* tp_richcompare */
1840 0, /* tp_weaklistoffset */
1841 PyObject_SelfIter, /* tp_iter */
1842 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001843 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001844 0, /* tp_members */
1845 0, /* tp_getset */
1846 0, /* tp_base */
1847 0, /* tp_dict */
1848 0, /* tp_descr_get */
1849 0, /* tp_descr_set */
1850 0, /* tp_dictoffset */
1851 0, /* tp_init */
1852 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001853 itertools_starmap, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001854 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001855};
1856
1857
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001858/* chain object **************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001859
1860typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001861 PyObject_HEAD
1862 PyObject *source; /* Iterator over input iterables */
1863 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001864} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001865
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001866static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001867
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001868static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00001869chain_new_internal(PyTypeObject *type, PyObject *source)
1870{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001871 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001872
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001873 lz = (chainobject *)type->tp_alloc(type, 0);
1874 if (lz == NULL) {
1875 Py_DECREF(source);
1876 return NULL;
1877 }
1878
1879 lz->source = source;
1880 lz->active = NULL;
1881 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00001882}
1883
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001884static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001885chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001886{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001887 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001888
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001889 if (type == &chain_type && !_PyArg_NoKeywords("chain", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001890 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001891
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001892 source = PyObject_GetIter(args);
1893 if (source == NULL)
1894 return NULL;
1895
1896 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00001897}
1898
Tal Einatc4bccd32018-09-12 00:49:13 +03001899/*[clinic input]
1900@classmethod
1901itertools.chain.from_iterable
1902 iterable as arg: object
1903 /
1904Alternative chain() constructor taking a single iterable argument that evaluates lazily.
1905[clinic start generated code]*/
1906
Christian Heimesf16baeb2008-02-29 14:57:44 +00001907static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001908itertools_chain_from_iterable(PyTypeObject *type, PyObject *arg)
1909/*[clinic end generated code: output=667ae7a7f7b68654 input=72c39e3a2ca3be85]*/
Christian Heimesf16baeb2008-02-29 14:57:44 +00001910{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001911 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001912
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001913 source = PyObject_GetIter(arg);
1914 if (source == NULL)
1915 return NULL;
1916
1917 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001918}
1919
1920static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001921chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001922{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001923 PyObject_GC_UnTrack(lz);
1924 Py_XDECREF(lz->active);
1925 Py_XDECREF(lz->source);
1926 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001927}
1928
Raymond Hettinger2012f172003-02-07 05:32:58 +00001929static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001930chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00001931{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001932 Py_VISIT(lz->source);
1933 Py_VISIT(lz->active);
1934 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00001935}
1936
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001937static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001938chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001939{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001940 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001941
T. Wouters5466d4a2017-03-30 09:58:35 -07001942 /* lz->source is the iterator of iterables. If it's NULL, we've already
1943 * consumed them all. lz->active is the current iterator. If it's NULL,
1944 * we should grab a new one from lz->source. */
1945 while (lz->source != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001946 if (lz->active == NULL) {
T. Wouters5466d4a2017-03-30 09:58:35 -07001947 PyObject *iterable = PyIter_Next(lz->source);
1948 if (iterable == NULL) {
1949 Py_CLEAR(lz->source);
1950 return NULL; /* no more input sources */
1951 }
1952 lz->active = PyObject_GetIter(iterable);
1953 Py_DECREF(iterable);
1954 if (lz->active == NULL) {
1955 Py_CLEAR(lz->source);
1956 return NULL; /* input not iterable */
1957 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001958 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001959 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
1960 if (item != NULL)
1961 return item;
1962 if (PyErr_Occurred()) {
1963 if (PyErr_ExceptionMatches(PyExc_StopIteration))
1964 PyErr_Clear();
1965 else
1966 return NULL; /* input raised an exception */
1967 }
1968 /* lz->active is consumed, try with the next iterable. */
1969 Py_CLEAR(lz->active);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001970 }
T. Wouters5466d4a2017-03-30 09:58:35 -07001971 /* Everything had been consumed already. */
1972 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001973}
1974
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001975static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301976chain_reduce(chainobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001977{
1978 if (lz->source) {
1979 /* we can't pickle function objects (itertools.from_iterable) so
1980 * we must use setstate to replace the iterable. One day we
1981 * will fix pickling of functions
1982 */
1983 if (lz->active) {
1984 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
1985 } else {
1986 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
1987 }
1988 } else {
1989 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
1990 }
1991 return NULL;
1992}
1993
1994static PyObject *
1995chain_setstate(chainobject *lz, PyObject *state)
1996{
1997 PyObject *source, *active=NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001998
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001999 if (!PyTuple_Check(state)) {
2000 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002001 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002002 }
2003 if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
2004 return NULL;
2005 }
2006 if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
2007 PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
2008 return NULL;
2009 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002010
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002011 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002012 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002013 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002014 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002015 Py_RETURN_NONE;
2016}
2017
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002018PyDoc_STRVAR(chain_doc,
2019"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002020\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002021Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002022first iterable until it is exhausted, then elements from the next\n\
2023iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002024
Christian Heimesf16baeb2008-02-29 14:57:44 +00002025static PyMethodDef chain_methods[] = {
Tal Einatc4bccd32018-09-12 00:49:13 +03002026 ITERTOOLS_CHAIN_FROM_ITERABLE_METHODDEF
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002027 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
2028 reduce_doc},
2029 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
2030 setstate_doc},
Ethan Smitha8403d02020-04-09 20:28:08 -07002031 {"__class_getitem__", (PyCFunction)Py_GenericAlias,
2032 METH_O|METH_CLASS, PyDoc_STR("See PEP 585")},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07002033 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00002034};
2035
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002036static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002037 PyVarObject_HEAD_INIT(NULL, 0)
2038 "itertools.chain", /* tp_name */
2039 sizeof(chainobject), /* tp_basicsize */
2040 0, /* tp_itemsize */
2041 /* methods */
2042 (destructor)chain_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002043 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002044 0, /* tp_getattr */
2045 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002046 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002047 0, /* tp_repr */
2048 0, /* tp_as_number */
2049 0, /* tp_as_sequence */
2050 0, /* tp_as_mapping */
2051 0, /* tp_hash */
2052 0, /* tp_call */
2053 0, /* tp_str */
2054 PyObject_GenericGetAttr, /* tp_getattro */
2055 0, /* tp_setattro */
2056 0, /* tp_as_buffer */
2057 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2058 Py_TPFLAGS_BASETYPE, /* tp_flags */
2059 chain_doc, /* tp_doc */
2060 (traverseproc)chain_traverse, /* tp_traverse */
2061 0, /* tp_clear */
2062 0, /* tp_richcompare */
2063 0, /* tp_weaklistoffset */
2064 PyObject_SelfIter, /* tp_iter */
2065 (iternextfunc)chain_next, /* tp_iternext */
2066 chain_methods, /* tp_methods */
2067 0, /* tp_members */
2068 0, /* tp_getset */
2069 0, /* tp_base */
2070 0, /* tp_dict */
2071 0, /* tp_descr_get */
2072 0, /* tp_descr_set */
2073 0, /* tp_dictoffset */
2074 0, /* tp_init */
2075 0, /* tp_alloc */
2076 chain_new, /* tp_new */
2077 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002078};
2079
2080
Christian Heimesc3f30c42008-02-22 16:37:40 +00002081/* product object ************************************************************/
2082
2083typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002084 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002085 PyObject *pools; /* tuple of pool tuples */
2086 Py_ssize_t *indices; /* one index per pool */
2087 PyObject *result; /* most recently returned result tuple */
2088 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002089} productobject;
2090
2091static PyTypeObject product_type;
2092
2093static PyObject *
2094product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2095{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002096 productobject *lz;
2097 Py_ssize_t nargs, npools, repeat=1;
2098 PyObject *pools = NULL;
2099 Py_ssize_t *indices = NULL;
2100 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002101
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002102 if (kwds != NULL) {
2103 char *kwlist[] = {"repeat", 0};
2104 PyObject *tmpargs = PyTuple_New(0);
2105 if (tmpargs == NULL)
2106 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002107 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2108 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002109 Py_DECREF(tmpargs);
2110 return NULL;
2111 }
2112 Py_DECREF(tmpargs);
2113 if (repeat < 0) {
2114 PyErr_SetString(PyExc_ValueError,
2115 "repeat argument cannot be negative");
2116 return NULL;
2117 }
2118 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002119
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002120 assert(PyTuple_CheckExact(args));
2121 if (repeat == 0) {
2122 nargs = 0;
2123 } else {
2124 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002125 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002126 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2127 return NULL;
2128 }
2129 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002130 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002131
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002132 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002133 if (indices == NULL) {
2134 PyErr_NoMemory();
2135 goto error;
2136 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002137
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002138 pools = PyTuple_New(npools);
2139 if (pools == NULL)
2140 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002141
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002142 for (i=0; i < nargs ; ++i) {
2143 PyObject *item = PyTuple_GET_ITEM(args, i);
2144 PyObject *pool = PySequence_Tuple(item);
2145 if (pool == NULL)
2146 goto error;
2147 PyTuple_SET_ITEM(pools, i, pool);
2148 indices[i] = 0;
2149 }
2150 for ( ; i < npools; ++i) {
2151 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2152 Py_INCREF(pool);
2153 PyTuple_SET_ITEM(pools, i, pool);
2154 indices[i] = 0;
2155 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002156
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002157 /* create productobject structure */
2158 lz = (productobject *)type->tp_alloc(type, 0);
2159 if (lz == NULL)
2160 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002161
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002162 lz->pools = pools;
2163 lz->indices = indices;
2164 lz->result = NULL;
2165 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002166
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002167 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002168
2169error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002170 if (indices != NULL)
2171 PyMem_Free(indices);
2172 Py_XDECREF(pools);
2173 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002174}
2175
2176static void
2177product_dealloc(productobject *lz)
2178{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002179 PyObject_GC_UnTrack(lz);
2180 Py_XDECREF(lz->pools);
2181 Py_XDECREF(lz->result);
2182 if (lz->indices != NULL)
2183 PyMem_Free(lz->indices);
2184 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002185}
2186
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002187static PyObject *
2188product_sizeof(productobject *lz, void *unused)
2189{
2190 Py_ssize_t res;
2191
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002192 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002193 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2194 return PyLong_FromSsize_t(res);
2195}
2196
2197PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2198
Christian Heimesc3f30c42008-02-22 16:37:40 +00002199static int
2200product_traverse(productobject *lz, visitproc visit, void *arg)
2201{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002202 Py_VISIT(lz->pools);
2203 Py_VISIT(lz->result);
2204 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002205}
2206
2207static PyObject *
2208product_next(productobject *lz)
2209{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002210 PyObject *pool;
2211 PyObject *elem;
2212 PyObject *oldelem;
2213 PyObject *pools = lz->pools;
2214 PyObject *result = lz->result;
2215 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2216 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002217
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002218 if (lz->stopped)
2219 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002220
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002221 if (result == NULL) {
2222 /* On the first pass, return an initial tuple filled with the
2223 first element from each pool. */
2224 result = PyTuple_New(npools);
2225 if (result == NULL)
2226 goto empty;
2227 lz->result = result;
2228 for (i=0; i < npools; i++) {
2229 pool = PyTuple_GET_ITEM(pools, i);
2230 if (PyTuple_GET_SIZE(pool) == 0)
2231 goto empty;
2232 elem = PyTuple_GET_ITEM(pool, 0);
2233 Py_INCREF(elem);
2234 PyTuple_SET_ITEM(result, i, elem);
2235 }
2236 } else {
2237 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002238
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002239 /* Copy the previous result tuple or re-use it if available */
2240 if (Py_REFCNT(result) > 1) {
2241 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002242 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002243 if (result == NULL)
2244 goto empty;
2245 lz->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002246 Py_DECREF(old_result);
2247 }
2248 /* Now, we've got the only copy so we can update it in-place */
2249 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002250
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002251 /* Update the pool indices right-to-left. Only advance to the
2252 next pool when the previous one rolls-over */
2253 for (i=npools-1 ; i >= 0 ; i--) {
2254 pool = PyTuple_GET_ITEM(pools, i);
2255 indices[i]++;
2256 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2257 /* Roll-over and advance to next pool */
2258 indices[i] = 0;
2259 elem = PyTuple_GET_ITEM(pool, 0);
2260 Py_INCREF(elem);
2261 oldelem = PyTuple_GET_ITEM(result, i);
2262 PyTuple_SET_ITEM(result, i, elem);
2263 Py_DECREF(oldelem);
2264 } else {
2265 /* No rollover. Just increment and stop here. */
2266 elem = PyTuple_GET_ITEM(pool, indices[i]);
2267 Py_INCREF(elem);
2268 oldelem = PyTuple_GET_ITEM(result, i);
2269 PyTuple_SET_ITEM(result, i, elem);
2270 Py_DECREF(oldelem);
2271 break;
2272 }
2273 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002274
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002275 /* If i is negative, then the indices have all rolled-over
2276 and we're done. */
2277 if (i < 0)
2278 goto empty;
2279 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002280
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002281 Py_INCREF(result);
2282 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002283
2284empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002285 lz->stopped = 1;
2286 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002287}
2288
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002289static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302290product_reduce(productobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002291{
2292 if (lz->stopped) {
2293 return Py_BuildValue("O(())", Py_TYPE(lz));
2294 } else if (lz->result == NULL) {
2295 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2296 } else {
2297 PyObject *indices;
2298 Py_ssize_t n, i;
2299
2300 /* we must pickle the indices use them for setstate, and
2301 * additionally indicate that the iterator has started
2302 */
2303 n = PyTuple_GET_SIZE(lz->pools);
2304 indices = PyTuple_New(n);
2305 if (indices == NULL)
2306 return NULL;
2307 for (i=0; i<n; i++){
2308 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2309 if (!index) {
2310 Py_DECREF(indices);
2311 return NULL;
2312 }
2313 PyTuple_SET_ITEM(indices, i, index);
2314 }
2315 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2316 }
2317}
2318
2319static PyObject *
2320product_setstate(productobject *lz, PyObject *state)
2321{
2322 PyObject *result;
2323 Py_ssize_t n, i;
2324
2325 n = PyTuple_GET_SIZE(lz->pools);
2326 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2327 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2328 return NULL;
2329 }
2330 for (i=0; i<n; i++)
2331 {
2332 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2333 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002334 PyObject* pool;
2335 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002336 if (index < 0 && PyErr_Occurred())
2337 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002338 pool = PyTuple_GET_ITEM(lz->pools, i);
2339 poolsize = PyTuple_GET_SIZE(pool);
2340 if (poolsize == 0) {
2341 lz->stopped = 1;
2342 Py_RETURN_NONE;
2343 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002344 /* clamp the index */
2345 if (index < 0)
2346 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002347 else if (index > poolsize-1)
2348 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002349 lz->indices[i] = index;
2350 }
2351
2352 result = PyTuple_New(n);
2353 if (!result)
2354 return NULL;
2355 for (i=0; i<n; i++) {
2356 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2357 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2358 Py_INCREF(element);
2359 PyTuple_SET_ITEM(result, i, element);
2360 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002361 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002362 Py_RETURN_NONE;
2363}
2364
2365static PyMethodDef product_methods[] = {
2366 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2367 reduce_doc},
2368 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2369 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002370 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2371 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002372 {NULL, NULL} /* sentinel */
2373};
2374
Christian Heimesc3f30c42008-02-22 16:37:40 +00002375PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002376"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002377\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002378Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002379For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2380The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2381cycle in a manner similar to an odometer (with the rightmost element changing\n\
2382on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002383To compute the product of an iterable with itself, specify the number\n\
2384of repetitions with the optional repeat keyword argument. For example,\n\
2385product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002386product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2387product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2388
2389static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002390 PyVarObject_HEAD_INIT(NULL, 0)
2391 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002392 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002393 0, /* tp_itemsize */
2394 /* methods */
2395 (destructor)product_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002396 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002397 0, /* tp_getattr */
2398 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002399 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002400 0, /* tp_repr */
2401 0, /* tp_as_number */
2402 0, /* tp_as_sequence */
2403 0, /* tp_as_mapping */
2404 0, /* tp_hash */
2405 0, /* tp_call */
2406 0, /* tp_str */
2407 PyObject_GenericGetAttr, /* tp_getattro */
2408 0, /* tp_setattro */
2409 0, /* tp_as_buffer */
2410 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2411 Py_TPFLAGS_BASETYPE, /* tp_flags */
2412 product_doc, /* tp_doc */
2413 (traverseproc)product_traverse, /* tp_traverse */
2414 0, /* tp_clear */
2415 0, /* tp_richcompare */
2416 0, /* tp_weaklistoffset */
2417 PyObject_SelfIter, /* tp_iter */
2418 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002419 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002420 0, /* tp_members */
2421 0, /* tp_getset */
2422 0, /* tp_base */
2423 0, /* tp_dict */
2424 0, /* tp_descr_get */
2425 0, /* tp_descr_set */
2426 0, /* tp_dictoffset */
2427 0, /* tp_init */
2428 0, /* tp_alloc */
2429 product_new, /* tp_new */
2430 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002431};
2432
2433
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002434/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002435
2436typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002437 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002438 PyObject *pool; /* input converted to a tuple */
2439 Py_ssize_t *indices; /* one index per result element */
2440 PyObject *result; /* most recently returned result tuple */
2441 Py_ssize_t r; /* size of result tuple */
2442 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002443} combinationsobject;
2444
Tal Einatc4bccd32018-09-12 00:49:13 +03002445
2446/*[clinic input]
2447@classmethod
2448itertools.combinations.__new__
2449 iterable: object
2450 r: Py_ssize_t
2451Return successive r-length combinations of elements in the iterable.
2452
2453combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)
2454[clinic start generated code]*/
2455
Christian Heimes380f7f22008-02-28 11:19:05 +00002456static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002457itertools_combinations_impl(PyTypeObject *type, PyObject *iterable,
2458 Py_ssize_t r)
2459/*[clinic end generated code: output=87a689b39c40039c input=06bede09e3da20f8]*/
Christian Heimes380f7f22008-02-28 11:19:05 +00002460{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002461 combinationsobject *co;
2462 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002463 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002464 Py_ssize_t *indices = NULL;
2465 Py_ssize_t i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002466
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002467 pool = PySequence_Tuple(iterable);
2468 if (pool == NULL)
2469 goto error;
2470 n = PyTuple_GET_SIZE(pool);
2471 if (r < 0) {
2472 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2473 goto error;
2474 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002475
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002476 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002477 if (indices == NULL) {
2478 PyErr_NoMemory();
2479 goto error;
2480 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002481
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002482 for (i=0 ; i<r ; i++)
2483 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002484
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002485 /* create combinationsobject structure */
2486 co = (combinationsobject *)type->tp_alloc(type, 0);
2487 if (co == NULL)
2488 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002489
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002490 co->pool = pool;
2491 co->indices = indices;
2492 co->result = NULL;
2493 co->r = r;
2494 co->stopped = r > n ? 1 : 0;
2495
2496 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002497
2498error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002499 if (indices != NULL)
2500 PyMem_Free(indices);
2501 Py_XDECREF(pool);
2502 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002503}
2504
2505static void
2506combinations_dealloc(combinationsobject *co)
2507{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002508 PyObject_GC_UnTrack(co);
2509 Py_XDECREF(co->pool);
2510 Py_XDECREF(co->result);
2511 if (co->indices != NULL)
2512 PyMem_Free(co->indices);
2513 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002514}
2515
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002516static PyObject *
2517combinations_sizeof(combinationsobject *co, void *unused)
2518{
2519 Py_ssize_t res;
2520
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002521 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002522 res += co->r * sizeof(Py_ssize_t);
2523 return PyLong_FromSsize_t(res);
2524}
2525
Christian Heimes380f7f22008-02-28 11:19:05 +00002526static int
2527combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2528{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002529 Py_VISIT(co->pool);
2530 Py_VISIT(co->result);
2531 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002532}
2533
2534static PyObject *
2535combinations_next(combinationsobject *co)
2536{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002537 PyObject *elem;
2538 PyObject *oldelem;
2539 PyObject *pool = co->pool;
2540 Py_ssize_t *indices = co->indices;
2541 PyObject *result = co->result;
2542 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2543 Py_ssize_t r = co->r;
2544 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002545
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002546 if (co->stopped)
2547 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002548
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002549 if (result == NULL) {
2550 /* On the first pass, initialize result tuple using the indices */
2551 result = PyTuple_New(r);
2552 if (result == NULL)
2553 goto empty;
2554 co->result = result;
2555 for (i=0; i<r ; i++) {
2556 index = indices[i];
2557 elem = PyTuple_GET_ITEM(pool, index);
2558 Py_INCREF(elem);
2559 PyTuple_SET_ITEM(result, i, elem);
2560 }
2561 } else {
2562 /* Copy the previous result tuple or re-use it if available */
2563 if (Py_REFCNT(result) > 1) {
2564 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002565 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002566 if (result == NULL)
2567 goto empty;
2568 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002569 Py_DECREF(old_result);
2570 }
2571 /* Now, we've got the only copy so we can update it in-place
2572 * CPython's empty tuple is a singleton and cached in
2573 * PyTuple's freelist.
2574 */
2575 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002576
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002577 /* Scan indices right-to-left until finding one that is not
2578 at its maximum (i + n - r). */
2579 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2580 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002581
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002582 /* If i is negative, then the indices are all at
2583 their maximum value and we're done. */
2584 if (i < 0)
2585 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002586
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002587 /* Increment the current index which we know is not at its
2588 maximum. Then move back to the right setting each index
2589 to its lowest possible value (one higher than the index
2590 to its left -- this maintains the sort order invariant). */
2591 indices[i]++;
2592 for (j=i+1 ; j<r ; j++)
2593 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002594
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002595 /* Update the result tuple for the new indices
2596 starting with i, the leftmost index that changed */
2597 for ( ; i<r ; i++) {
2598 index = indices[i];
2599 elem = PyTuple_GET_ITEM(pool, index);
2600 Py_INCREF(elem);
2601 oldelem = PyTuple_GET_ITEM(result, i);
2602 PyTuple_SET_ITEM(result, i, elem);
2603 Py_DECREF(oldelem);
2604 }
2605 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002606
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002607 Py_INCREF(result);
2608 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002609
2610empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002611 co->stopped = 1;
2612 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002613}
2614
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002615static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302616combinations_reduce(combinationsobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002617{
2618 if (lz->result == NULL) {
2619 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2620 } else if (lz->stopped) {
2621 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2622 } else {
2623 PyObject *indices;
2624 Py_ssize_t i;
2625
2626 /* we must pickle the indices and use them for setstate */
2627 indices = PyTuple_New(lz->r);
2628 if (!indices)
2629 return NULL;
2630 for (i=0; i<lz->r; i++)
2631 {
2632 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2633 if (!index) {
2634 Py_DECREF(indices);
2635 return NULL;
2636 }
2637 PyTuple_SET_ITEM(indices, i, index);
2638 }
2639
2640 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2641 }
2642}
2643
2644static PyObject *
2645combinations_setstate(combinationsobject *lz, PyObject *state)
2646{
2647 PyObject *result;
2648 Py_ssize_t i;
2649 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2650
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002651 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002652 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2653 return NULL;
2654 }
2655
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002656 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002657 Py_ssize_t max;
2658 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2659 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002660
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002661 if (index == -1 && PyErr_Occurred())
2662 return NULL; /* not an integer */
2663 max = i + n - lz->r;
2664 /* clamp the index (beware of negative max) */
2665 if (index > max)
2666 index = max;
2667 if (index < 0)
2668 index = 0;
2669 lz->indices[i] = index;
2670 }
2671
2672 result = PyTuple_New(lz->r);
2673 if (result == NULL)
2674 return NULL;
2675 for (i=0; i<lz->r; i++) {
2676 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2677 Py_INCREF(element);
2678 PyTuple_SET_ITEM(result, i, element);
2679 }
2680
Serhiy Storchaka48842712016-04-06 09:45:48 +03002681 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002682 Py_RETURN_NONE;
2683}
2684
2685static PyMethodDef combinations_methods[] = {
2686 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2687 reduce_doc},
2688 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2689 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002690 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2691 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002692 {NULL, NULL} /* sentinel */
2693};
2694
Christian Heimes380f7f22008-02-28 11:19:05 +00002695static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002696 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002697 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002698 sizeof(combinationsobject), /* tp_basicsize */
2699 0, /* tp_itemsize */
2700 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002701 (destructor)combinations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002702 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002703 0, /* tp_getattr */
2704 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002705 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002706 0, /* tp_repr */
2707 0, /* tp_as_number */
2708 0, /* tp_as_sequence */
2709 0, /* tp_as_mapping */
2710 0, /* tp_hash */
2711 0, /* tp_call */
2712 0, /* tp_str */
2713 PyObject_GenericGetAttr, /* tp_getattro */
2714 0, /* tp_setattro */
2715 0, /* tp_as_buffer */
2716 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2717 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03002718 itertools_combinations__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002719 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002720 0, /* tp_clear */
2721 0, /* tp_richcompare */
2722 0, /* tp_weaklistoffset */
2723 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002724 (iternextfunc)combinations_next, /* tp_iternext */
2725 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002726 0, /* tp_members */
2727 0, /* tp_getset */
2728 0, /* tp_base */
2729 0, /* tp_dict */
2730 0, /* tp_descr_get */
2731 0, /* tp_descr_set */
2732 0, /* tp_dictoffset */
2733 0, /* tp_init */
2734 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03002735 itertools_combinations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002736 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002737};
2738
2739
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002740/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002741
2742/* Equivalent to:
2743
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002744 def combinations_with_replacement(iterable, r):
2745 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2746 # number items returned: (n+r-1)! / r! / (n-1)!
2747 pool = tuple(iterable)
2748 n = len(pool)
2749 indices = [0] * r
2750 yield tuple(pool[i] for i in indices)
2751 while 1:
2752 for i in reversed(range(r)):
2753 if indices[i] != n - 1:
2754 break
2755 else:
2756 return
2757 indices[i:] = [indices[i] + 1] * (r - i)
2758 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002759
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002760 def combinations_with_replacement2(iterable, r):
2761 'Alternate version that filters from product()'
2762 pool = tuple(iterable)
2763 n = len(pool)
2764 for indices in product(range(n), repeat=r):
2765 if sorted(indices) == list(indices):
2766 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002767*/
2768typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002769 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002770 PyObject *pool; /* input converted to a tuple */
2771 Py_ssize_t *indices; /* one index per result element */
2772 PyObject *result; /* most recently returned result tuple */
2773 Py_ssize_t r; /* size of result tuple */
2774 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002775} cwrobject;
2776
Tal Einatc4bccd32018-09-12 00:49:13 +03002777/*[clinic input]
2778@classmethod
2779itertools.combinations_with_replacement.__new__
2780 iterable: object
2781 r: Py_ssize_t
2782Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats.
2783
2784combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2785[clinic start generated code]*/
2786
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002787static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002788itertools_combinations_with_replacement_impl(PyTypeObject *type,
2789 PyObject *iterable,
2790 Py_ssize_t r)
2791/*[clinic end generated code: output=48b26856d4e659ca input=dc2a8c7ba785fad7]*/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002792{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002793 cwrobject *co;
2794 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002795 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002796 Py_ssize_t *indices = NULL;
2797 Py_ssize_t i;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002798
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002799 pool = PySequence_Tuple(iterable);
2800 if (pool == NULL)
2801 goto error;
2802 n = PyTuple_GET_SIZE(pool);
2803 if (r < 0) {
2804 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2805 goto error;
2806 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002807
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002808 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002809 if (indices == NULL) {
2810 PyErr_NoMemory();
2811 goto error;
2812 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002813
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002814 for (i=0 ; i<r ; i++)
2815 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002816
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002817 /* create cwrobject structure */
2818 co = (cwrobject *)type->tp_alloc(type, 0);
2819 if (co == NULL)
2820 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002821
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002822 co->pool = pool;
2823 co->indices = indices;
2824 co->result = NULL;
2825 co->r = r;
2826 co->stopped = !n && r;
2827
2828 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002829
2830error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002831 if (indices != NULL)
2832 PyMem_Free(indices);
2833 Py_XDECREF(pool);
2834 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002835}
2836
2837static void
2838cwr_dealloc(cwrobject *co)
2839{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002840 PyObject_GC_UnTrack(co);
2841 Py_XDECREF(co->pool);
2842 Py_XDECREF(co->result);
2843 if (co->indices != NULL)
2844 PyMem_Free(co->indices);
2845 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002846}
2847
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002848static PyObject *
2849cwr_sizeof(cwrobject *co, void *unused)
2850{
2851 Py_ssize_t res;
2852
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002853 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002854 res += co->r * sizeof(Py_ssize_t);
2855 return PyLong_FromSsize_t(res);
2856}
2857
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002858static int
2859cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2860{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002861 Py_VISIT(co->pool);
2862 Py_VISIT(co->result);
2863 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002864}
2865
2866static PyObject *
2867cwr_next(cwrobject *co)
2868{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002869 PyObject *elem;
2870 PyObject *oldelem;
2871 PyObject *pool = co->pool;
2872 Py_ssize_t *indices = co->indices;
2873 PyObject *result = co->result;
2874 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2875 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05002876 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002877
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002878 if (co->stopped)
2879 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002880
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002881 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05002882 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002883 result = PyTuple_New(r);
2884 if (result == NULL)
2885 goto empty;
2886 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07002887 if (n > 0) {
2888 elem = PyTuple_GET_ITEM(pool, 0);
2889 for (i=0; i<r ; i++) {
2890 assert(indices[i] == 0);
2891 Py_INCREF(elem);
2892 PyTuple_SET_ITEM(result, i, elem);
2893 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002894 }
2895 } else {
2896 /* Copy the previous result tuple or re-use it if available */
2897 if (Py_REFCNT(result) > 1) {
2898 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002899 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002900 if (result == NULL)
2901 goto empty;
2902 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002903 Py_DECREF(old_result);
2904 }
2905 /* Now, we've got the only copy so we can update it in-place CPython's
2906 empty tuple is a singleton and cached in PyTuple's freelist. */
2907 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002908
Tim Peters9edb1682013-09-03 11:49:31 -05002909 /* Scan indices right-to-left until finding one that is not
2910 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002911 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2912 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002913
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002914 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05002915 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002916 if (i < 0)
2917 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002918
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002919 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05002920 maximum. Then set all to the right to the same value. */
2921 index = indices[i] + 1;
2922 assert(index < n);
2923 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002924 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05002925 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002926 Py_INCREF(elem);
2927 oldelem = PyTuple_GET_ITEM(result, i);
2928 PyTuple_SET_ITEM(result, i, elem);
2929 Py_DECREF(oldelem);
2930 }
2931 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002932
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002933 Py_INCREF(result);
2934 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002935
2936empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002937 co->stopped = 1;
2938 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002939}
2940
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002941static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302942cwr_reduce(cwrobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002943{
2944 if (lz->result == NULL) {
2945 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2946 } else if (lz->stopped) {
2947 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2948 } else {
2949 PyObject *indices;
2950 Py_ssize_t i;
2951
2952 /* we must pickle the indices and use them for setstate */
2953 indices = PyTuple_New(lz->r);
2954 if (!indices)
2955 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002956 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002957 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2958 if (!index) {
2959 Py_DECREF(indices);
2960 return NULL;
2961 }
2962 PyTuple_SET_ITEM(indices, i, index);
2963 }
2964
2965 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2966 }
2967}
2968
2969static PyObject *
2970cwr_setstate(cwrobject *lz, PyObject *state)
2971{
2972 PyObject *result;
2973 Py_ssize_t n, i;
2974
2975 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
2976 {
2977 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2978 return NULL;
2979 }
2980
2981 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002982 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002983 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2984 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002985
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002986 if (index < 0 && PyErr_Occurred())
2987 return NULL; /* not an integer */
2988 /* clamp the index */
2989 if (index < 0)
2990 index = 0;
2991 else if (index > n-1)
2992 index = n-1;
2993 lz->indices[i] = index;
2994 }
2995 result = PyTuple_New(lz->r);
2996 if (result == NULL)
2997 return NULL;
2998 for (i=0; i<lz->r; i++) {
2999 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
3000 Py_INCREF(element);
3001 PyTuple_SET_ITEM(result, i, element);
3002 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003003 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003004 Py_RETURN_NONE;
3005}
3006
3007static PyMethodDef cwr_methods[] = {
3008 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
3009 reduce_doc},
3010 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
3011 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003012 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
3013 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003014 {NULL, NULL} /* sentinel */
3015};
3016
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003017static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003018 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003019 "itertools.combinations_with_replacement", /* tp_name */
3020 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003021 0, /* tp_itemsize */
3022 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003023 (destructor)cwr_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003024 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003025 0, /* tp_getattr */
3026 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003027 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003028 0, /* tp_repr */
3029 0, /* tp_as_number */
3030 0, /* tp_as_sequence */
3031 0, /* tp_as_mapping */
3032 0, /* tp_hash */
3033 0, /* tp_call */
3034 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003035 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003036 0, /* tp_setattro */
3037 0, /* tp_as_buffer */
3038 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003039 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003040 itertools_combinations_with_replacement__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003041 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003042 0, /* tp_clear */
3043 0, /* tp_richcompare */
3044 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003045 PyObject_SelfIter, /* tp_iter */
3046 (iternextfunc)cwr_next, /* tp_iternext */
3047 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003048 0, /* tp_members */
3049 0, /* tp_getset */
3050 0, /* tp_base */
3051 0, /* tp_dict */
3052 0, /* tp_descr_get */
3053 0, /* tp_descr_set */
3054 0, /* tp_dictoffset */
3055 0, /* tp_init */
3056 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003057 itertools_combinations_with_replacement, /* tp_new */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003058 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003059};
3060
3061
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003062/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003063
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003064def permutations(iterable, r=None):
Sergey0078a0c2019-10-29 08:10:24 +03003065 # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
3066 # permutations(range(3)) --> 012 021 102 120 201 210
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003067 pool = tuple(iterable)
3068 n = len(pool)
3069 r = n if r is None else r
Sergey0078a0c2019-10-29 08:10:24 +03003070 if r > n:
3071 return
3072 indices = list(range(n))
3073 cycles = list(range(n, n-r, -1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003074 yield tuple(pool[i] for i in indices[:r])
3075 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003076 for i in reversed(range(r)):
3077 cycles[i] -= 1
3078 if cycles[i] == 0:
3079 indices[i:] = indices[i+1:] + indices[i:i+1]
3080 cycles[i] = n - i
3081 else:
3082 j = cycles[i]
3083 indices[i], indices[-j] = indices[-j], indices[i]
3084 yield tuple(pool[i] for i in indices[:r])
3085 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003086 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003087 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003088*/
3089
3090typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003091 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003092 PyObject *pool; /* input converted to a tuple */
3093 Py_ssize_t *indices; /* one index per element in the pool */
3094 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3095 PyObject *result; /* most recently returned result tuple */
3096 Py_ssize_t r; /* size of result tuple */
3097 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003098} permutationsobject;
3099
Tal Einatc4bccd32018-09-12 00:49:13 +03003100/*[clinic input]
3101@classmethod
3102itertools.permutations.__new__
3103 iterable: object
3104 r as robj: object = None
3105Return successive r-length permutations of elements in the iterable.
3106
3107permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)
3108[clinic start generated code]*/
3109
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003110static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003111itertools_permutations_impl(PyTypeObject *type, PyObject *iterable,
3112 PyObject *robj)
3113/*[clinic end generated code: output=296a72fa76d620ea input=57d0170a4ac0ec7a]*/
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003114{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003115 permutationsobject *po;
3116 Py_ssize_t n;
3117 Py_ssize_t r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003118 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003119 Py_ssize_t *indices = NULL;
3120 Py_ssize_t *cycles = NULL;
3121 Py_ssize_t i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003122
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003123 pool = PySequence_Tuple(iterable);
3124 if (pool == NULL)
3125 goto error;
3126 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003127
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003128 r = n;
3129 if (robj != Py_None) {
3130 if (!PyLong_Check(robj)) {
3131 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3132 goto error;
3133 }
3134 r = PyLong_AsSsize_t(robj);
3135 if (r == -1 && PyErr_Occurred())
3136 goto error;
3137 }
3138 if (r < 0) {
3139 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3140 goto error;
3141 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003142
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003143 indices = PyMem_New(Py_ssize_t, n);
3144 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003145 if (indices == NULL || cycles == NULL) {
3146 PyErr_NoMemory();
3147 goto error;
3148 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003149
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003150 for (i=0 ; i<n ; i++)
3151 indices[i] = i;
3152 for (i=0 ; i<r ; i++)
3153 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003154
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003155 /* create permutationsobject structure */
3156 po = (permutationsobject *)type->tp_alloc(type, 0);
3157 if (po == NULL)
3158 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003159
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003160 po->pool = pool;
3161 po->indices = indices;
3162 po->cycles = cycles;
3163 po->result = NULL;
3164 po->r = r;
3165 po->stopped = r > n ? 1 : 0;
3166
3167 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003168
3169error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003170 if (indices != NULL)
3171 PyMem_Free(indices);
3172 if (cycles != NULL)
3173 PyMem_Free(cycles);
3174 Py_XDECREF(pool);
3175 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003176}
3177
3178static void
3179permutations_dealloc(permutationsobject *po)
3180{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003181 PyObject_GC_UnTrack(po);
3182 Py_XDECREF(po->pool);
3183 Py_XDECREF(po->result);
3184 PyMem_Free(po->indices);
3185 PyMem_Free(po->cycles);
3186 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003187}
3188
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003189static PyObject *
3190permutations_sizeof(permutationsobject *po, void *unused)
3191{
3192 Py_ssize_t res;
3193
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003194 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003195 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3196 res += po->r * sizeof(Py_ssize_t);
3197 return PyLong_FromSsize_t(res);
3198}
3199
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003200static int
3201permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3202{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003203 Py_VISIT(po->pool);
3204 Py_VISIT(po->result);
3205 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003206}
3207
3208static PyObject *
3209permutations_next(permutationsobject *po)
3210{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003211 PyObject *elem;
3212 PyObject *oldelem;
3213 PyObject *pool = po->pool;
3214 Py_ssize_t *indices = po->indices;
3215 Py_ssize_t *cycles = po->cycles;
3216 PyObject *result = po->result;
3217 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3218 Py_ssize_t r = po->r;
3219 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003220
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003221 if (po->stopped)
3222 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003223
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003224 if (result == NULL) {
3225 /* On the first pass, initialize result tuple using the indices */
3226 result = PyTuple_New(r);
3227 if (result == NULL)
3228 goto empty;
3229 po->result = result;
3230 for (i=0; i<r ; i++) {
3231 index = indices[i];
3232 elem = PyTuple_GET_ITEM(pool, index);
3233 Py_INCREF(elem);
3234 PyTuple_SET_ITEM(result, i, elem);
3235 }
3236 } else {
3237 if (n == 0)
3238 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003239
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003240 /* Copy the previous result tuple or re-use it if available */
3241 if (Py_REFCNT(result) > 1) {
3242 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05003243 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003244 if (result == NULL)
3245 goto empty;
3246 po->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003247 Py_DECREF(old_result);
3248 }
3249 /* Now, we've got the only copy so we can update it in-place */
3250 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003251
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003252 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3253 for (i=r-1 ; i>=0 ; i--) {
3254 cycles[i] -= 1;
3255 if (cycles[i] == 0) {
3256 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3257 index = indices[i];
3258 for (j=i ; j<n-1 ; j++)
3259 indices[j] = indices[j+1];
3260 indices[n-1] = index;
3261 cycles[i] = n - i;
3262 } else {
3263 j = cycles[i];
3264 index = indices[i];
3265 indices[i] = indices[n-j];
3266 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003267
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003268 for (k=i; k<r ; k++) {
3269 /* start with i, the leftmost element that changed */
3270 /* yield tuple(pool[k] for k in indices[:r]) */
3271 index = indices[k];
3272 elem = PyTuple_GET_ITEM(pool, index);
3273 Py_INCREF(elem);
3274 oldelem = PyTuple_GET_ITEM(result, k);
3275 PyTuple_SET_ITEM(result, k, elem);
3276 Py_DECREF(oldelem);
3277 }
3278 break;
3279 }
3280 }
3281 /* If i is negative, then the cycles have all
3282 rolled-over and we're done. */
3283 if (i < 0)
3284 goto empty;
3285 }
3286 Py_INCREF(result);
3287 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003288
3289empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003290 po->stopped = 1;
3291 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003292}
3293
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003294static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303295permutations_reduce(permutationsobject *po, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003296{
3297 if (po->result == NULL) {
3298 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3299 } else if (po->stopped) {
3300 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3301 } else {
3302 PyObject *indices=NULL, *cycles=NULL;
3303 Py_ssize_t n, i;
3304
3305 /* we must pickle the indices and cycles and use them for setstate */
3306 n = PyTuple_GET_SIZE(po->pool);
3307 indices = PyTuple_New(n);
3308 if (indices == NULL)
3309 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003310 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003311 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3312 if (!index)
3313 goto err;
3314 PyTuple_SET_ITEM(indices, i, index);
3315 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003316
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003317 cycles = PyTuple_New(po->r);
3318 if (cycles == NULL)
3319 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003320 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003321 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3322 if (!index)
3323 goto err;
3324 PyTuple_SET_ITEM(cycles, i, index);
3325 }
3326 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3327 po->pool, po->r,
3328 indices, cycles);
3329 err:
3330 Py_XDECREF(indices);
3331 Py_XDECREF(cycles);
3332 return NULL;
3333 }
3334}
3335
3336static PyObject *
3337permutations_setstate(permutationsobject *po, PyObject *state)
3338{
3339 PyObject *indices, *cycles, *result;
3340 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003341
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003342 if (!PyTuple_Check(state)) {
3343 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3344 return NULL;
3345 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003346 if (!PyArg_ParseTuple(state, "O!O!",
3347 &PyTuple_Type, &indices,
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003348 &PyTuple_Type, &cycles)) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003349 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003350 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003351
3352 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003353 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003354 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3355 return NULL;
3356 }
3357
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003358 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003359 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3360 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3361 if (index < 0 && PyErr_Occurred())
3362 return NULL; /* not an integer */
3363 /* clamp the index */
3364 if (index < 0)
3365 index = 0;
3366 else if (index > n-1)
3367 index = n-1;
3368 po->indices[i] = index;
3369 }
3370
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003371 for (i=0; i<po->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003372 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3373 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3374 if (index < 0 && PyErr_Occurred())
3375 return NULL; /* not an integer */
3376 if (index < 1)
3377 index = 1;
3378 else if (index > n-i)
3379 index = n-i;
3380 po->cycles[i] = index;
3381 }
3382 result = PyTuple_New(po->r);
3383 if (result == NULL)
3384 return NULL;
3385 for (i=0; i<po->r; i++) {
3386 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3387 Py_INCREF(element);
3388 PyTuple_SET_ITEM(result, i, element);
3389 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003390 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003391 Py_RETURN_NONE;
3392}
3393
3394static PyMethodDef permuations_methods[] = {
3395 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3396 reduce_doc},
3397 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3398 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003399 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3400 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003401 {NULL, NULL} /* sentinel */
3402};
3403
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003404static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003405 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003406 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003407 sizeof(permutationsobject), /* tp_basicsize */
3408 0, /* tp_itemsize */
3409 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003410 (destructor)permutations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003411 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003412 0, /* tp_getattr */
3413 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003414 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003415 0, /* tp_repr */
3416 0, /* tp_as_number */
3417 0, /* tp_as_sequence */
3418 0, /* tp_as_mapping */
3419 0, /* tp_hash */
3420 0, /* tp_call */
3421 0, /* tp_str */
3422 PyObject_GenericGetAttr, /* tp_getattro */
3423 0, /* tp_setattro */
3424 0, /* tp_as_buffer */
3425 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3426 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003427 itertools_permutations__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003428 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003429 0, /* tp_clear */
3430 0, /* tp_richcompare */
3431 0, /* tp_weaklistoffset */
3432 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003433 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003434 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003435 0, /* tp_members */
3436 0, /* tp_getset */
3437 0, /* tp_base */
3438 0, /* tp_dict */
3439 0, /* tp_descr_get */
3440 0, /* tp_descr_set */
3441 0, /* tp_dictoffset */
3442 0, /* tp_init */
3443 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003444 itertools_permutations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003445 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003446};
3447
Tal Einatc4bccd32018-09-12 00:49:13 +03003448
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003449/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003450
3451typedef struct {
3452 PyObject_HEAD
3453 PyObject *total;
3454 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003455 PyObject *binop;
Lisa Roach9718b592018-09-23 17:34:59 -07003456 PyObject *initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003457} accumulateobject;
3458
Tal Einatc4bccd32018-09-12 00:49:13 +03003459/*[clinic input]
3460@classmethod
3461itertools.accumulate.__new__
3462 iterable: object
3463 func as binop: object = None
Lisa Roach9718b592018-09-23 17:34:59 -07003464 *
3465 initial: object = None
Tal Einatc4bccd32018-09-12 00:49:13 +03003466Return series of accumulated sums (or other binary function results).
3467[clinic start generated code]*/
3468
Raymond Hettinger482ba772010-12-01 22:48:00 +00003469static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003470itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
Lisa Roach9718b592018-09-23 17:34:59 -07003471 PyObject *binop, PyObject *initial)
3472/*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003473{
Raymond Hettinger482ba772010-12-01 22:48:00 +00003474 PyObject *it;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003475 accumulateobject *lz;
3476
Raymond Hettinger482ba772010-12-01 22:48:00 +00003477 /* Get iterator. */
3478 it = PyObject_GetIter(iterable);
3479 if (it == NULL)
3480 return NULL;
3481
Raymond Hettinger482ba772010-12-01 22:48:00 +00003482 /* create accumulateobject structure */
3483 lz = (accumulateobject *)type->tp_alloc(type, 0);
3484 if (lz == NULL) {
3485 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003486 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003487 }
3488
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003489 if (binop != Py_None) {
3490 Py_XINCREF(binop);
3491 lz->binop = binop;
3492 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003493 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003494 lz->it = it;
Lisa Roach9718b592018-09-23 17:34:59 -07003495 Py_XINCREF(initial);
3496 lz->initial = initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003497 return (PyObject *)lz;
3498}
3499
3500static void
3501accumulate_dealloc(accumulateobject *lz)
3502{
3503 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003504 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003505 Py_XDECREF(lz->total);
3506 Py_XDECREF(lz->it);
Lisa Roach9718b592018-09-23 17:34:59 -07003507 Py_XDECREF(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003508 Py_TYPE(lz)->tp_free(lz);
3509}
3510
3511static int
3512accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3513{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003514 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003515 Py_VISIT(lz->it);
3516 Py_VISIT(lz->total);
Lisa Roach9718b592018-09-23 17:34:59 -07003517 Py_VISIT(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003518 return 0;
3519}
3520
3521static PyObject *
3522accumulate_next(accumulateobject *lz)
3523{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003524 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003525
Lisa Roach9718b592018-09-23 17:34:59 -07003526 if (lz->initial != Py_None) {
3527 lz->total = lz->initial;
3528 Py_INCREF(Py_None);
3529 lz->initial = Py_None;
3530 Py_INCREF(lz->total);
3531 return lz->total;
3532 }
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003533 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003534 if (val == NULL)
3535 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003536
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003537 if (lz->total == NULL) {
3538 Py_INCREF(val);
3539 lz->total = val;
3540 return lz->total;
3541 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003542
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003543 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003544 newtotal = PyNumber_Add(lz->total, val);
3545 else
3546 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003547 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003548 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003549 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003550
Raymond Hettinger482ba772010-12-01 22:48:00 +00003551 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003552 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003553 return newtotal;
3554}
3555
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003556static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303557accumulate_reduce(accumulateobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003558{
Lisa Roach9718b592018-09-23 17:34:59 -07003559 if (lz->initial != Py_None) {
3560 PyObject *it;
3561
3562 assert(lz->total == NULL);
3563 if (PyType_Ready(&chain_type) < 0)
3564 return NULL;
3565 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3566 lz->initial, lz->it);
3567 if (it == NULL)
3568 return NULL;
3569 return Py_BuildValue("O(NO)O", Py_TYPE(lz),
3570 it, lz->binop?lz->binop:Py_None, Py_None);
3571 }
Serhiy Storchakad5516252016-03-06 14:00:45 +02003572 if (lz->total == Py_None) {
3573 PyObject *it;
3574
3575 if (PyType_Ready(&chain_type) < 0)
3576 return NULL;
3577 if (PyType_Ready(&islice_type) < 0)
3578 return NULL;
3579 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3580 lz->total, lz->it);
3581 if (it == NULL)
3582 return NULL;
3583 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3584 it, lz->binop ? lz->binop : Py_None);
3585 if (it == NULL)
3586 return NULL;
3587 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3588 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003589 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3590 lz->it, lz->binop?lz->binop:Py_None,
3591 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003592}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003593
3594static PyObject *
3595accumulate_setstate(accumulateobject *lz, PyObject *state)
3596{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003597 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003598 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003599 Py_RETURN_NONE;
3600}
3601
3602static PyMethodDef accumulate_methods[] = {
3603 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3604 reduce_doc},
3605 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3606 setstate_doc},
3607 {NULL, NULL} /* sentinel */
3608};
3609
Raymond Hettinger482ba772010-12-01 22:48:00 +00003610static PyTypeObject accumulate_type = {
3611 PyVarObject_HEAD_INIT(NULL, 0)
3612 "itertools.accumulate", /* tp_name */
3613 sizeof(accumulateobject), /* tp_basicsize */
3614 0, /* tp_itemsize */
3615 /* methods */
3616 (destructor)accumulate_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003617 0, /* tp_vectorcall_offset */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003618 0, /* tp_getattr */
3619 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003620 0, /* tp_as_async */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003621 0, /* tp_repr */
3622 0, /* tp_as_number */
3623 0, /* tp_as_sequence */
3624 0, /* tp_as_mapping */
3625 0, /* tp_hash */
3626 0, /* tp_call */
3627 0, /* tp_str */
3628 PyObject_GenericGetAttr, /* tp_getattro */
3629 0, /* tp_setattro */
3630 0, /* tp_as_buffer */
3631 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3632 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003633 itertools_accumulate__doc__, /* tp_doc */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003634 (traverseproc)accumulate_traverse, /* tp_traverse */
3635 0, /* tp_clear */
3636 0, /* tp_richcompare */
3637 0, /* tp_weaklistoffset */
3638 PyObject_SelfIter, /* tp_iter */
3639 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003640 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003641 0, /* tp_members */
3642 0, /* tp_getset */
3643 0, /* tp_base */
3644 0, /* tp_dict */
3645 0, /* tp_descr_get */
3646 0, /* tp_descr_set */
3647 0, /* tp_dictoffset */
3648 0, /* tp_init */
3649 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003650 itertools_accumulate, /* tp_new */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003651 PyObject_GC_Del, /* tp_free */
3652};
3653
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003654
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003655/* compress object ************************************************************/
3656
3657/* Equivalent to:
3658
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003659 def compress(data, selectors):
3660 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3661 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003662*/
3663
3664typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003665 PyObject_HEAD
3666 PyObject *data;
3667 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003668} compressobject;
3669
Tal Einatc4bccd32018-09-12 00:49:13 +03003670/*[clinic input]
3671@classmethod
3672itertools.compress.__new__
3673 data as seq1: object
3674 selectors as seq2: object
3675Return data elements corresponding to true selector elements.
3676
3677Forms a shorter iterator from selected data elements using the selectors to
3678choose the data elements.
3679[clinic start generated code]*/
3680
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003681static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003682itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
3683/*[clinic end generated code: output=7e67157212ed09e0 input=79596d7cd20c77e5]*/
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003684{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003685 PyObject *data=NULL, *selectors=NULL;
3686 compressobject *lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003687
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003688 data = PyObject_GetIter(seq1);
3689 if (data == NULL)
3690 goto fail;
3691 selectors = PyObject_GetIter(seq2);
3692 if (selectors == NULL)
3693 goto fail;
3694
3695 /* create compressobject structure */
3696 lz = (compressobject *)type->tp_alloc(type, 0);
3697 if (lz == NULL)
3698 goto fail;
3699 lz->data = data;
3700 lz->selectors = selectors;
3701 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003702
3703fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003704 Py_XDECREF(data);
3705 Py_XDECREF(selectors);
3706 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003707}
3708
3709static void
3710compress_dealloc(compressobject *lz)
3711{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003712 PyObject_GC_UnTrack(lz);
3713 Py_XDECREF(lz->data);
3714 Py_XDECREF(lz->selectors);
3715 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003716}
3717
3718static int
3719compress_traverse(compressobject *lz, visitproc visit, void *arg)
3720{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003721 Py_VISIT(lz->data);
3722 Py_VISIT(lz->selectors);
3723 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003724}
3725
3726static PyObject *
3727compress_next(compressobject *lz)
3728{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003729 PyObject *data = lz->data, *selectors = lz->selectors;
3730 PyObject *datum, *selector;
3731 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3732 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3733 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003734
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003735 while (1) {
3736 /* Steps: get datum, get selector, evaluate selector.
3737 Order is important (to match the pure python version
3738 in terms of which input gets a chance to raise an
3739 exception first).
3740 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003741
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003742 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003743 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003744 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003745
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003746 selector = selectornext(selectors);
3747 if (selector == NULL) {
3748 Py_DECREF(datum);
3749 return NULL;
3750 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003751
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003752 ok = PyObject_IsTrue(selector);
3753 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003754 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003755 return datum;
3756 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003757 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003758 return NULL;
3759 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003760}
3761
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003762static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303763compress_reduce(compressobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003764{
3765 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3766 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003767}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003768
3769static PyMethodDef compress_methods[] = {
3770 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3771 reduce_doc},
3772 {NULL, NULL} /* sentinel */
3773};
3774
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003775static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003776 PyVarObject_HEAD_INIT(NULL, 0)
3777 "itertools.compress", /* tp_name */
3778 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003779 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003780 /* methods */
3781 (destructor)compress_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003782 0, /* tp_vectorcall_offset */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003783 0, /* tp_getattr */
3784 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003785 0, /* tp_as_async */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003786 0, /* tp_repr */
3787 0, /* tp_as_number */
3788 0, /* tp_as_sequence */
3789 0, /* tp_as_mapping */
3790 0, /* tp_hash */
3791 0, /* tp_call */
3792 0, /* tp_str */
3793 PyObject_GenericGetAttr, /* tp_getattro */
3794 0, /* tp_setattro */
3795 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003796 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003797 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003798 itertools_compress__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003799 (traverseproc)compress_traverse, /* tp_traverse */
3800 0, /* tp_clear */
3801 0, /* tp_richcompare */
3802 0, /* tp_weaklistoffset */
3803 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003804 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003805 compress_methods, /* tp_methods */
3806 0, /* tp_members */
3807 0, /* tp_getset */
3808 0, /* tp_base */
3809 0, /* tp_dict */
3810 0, /* tp_descr_get */
3811 0, /* tp_descr_set */
3812 0, /* tp_dictoffset */
3813 0, /* tp_init */
3814 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003815 itertools_compress, /* tp_new */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003816 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003817};
3818
3819
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003820/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003821
3822typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003823 PyObject_HEAD
3824 PyObject *func;
3825 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003826} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003827
Tal Einatc4bccd32018-09-12 00:49:13 +03003828/*[clinic input]
3829@classmethod
3830itertools.filterfalse.__new__
3831 function as func: object
3832 iterable as seq: object
3833 /
3834Return those items of iterable for which function(item) is false.
3835
3836If function is None, return the items that are false.
3837[clinic start generated code]*/
3838
Raymond Hettinger60eca932003-02-09 06:40:58 +00003839static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003840itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
3841/*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003842{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003843 PyObject *it;
3844 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003845
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003846 /* Get iterator. */
3847 it = PyObject_GetIter(seq);
3848 if (it == NULL)
3849 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003850
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003851 /* create filterfalseobject structure */
3852 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3853 if (lz == NULL) {
3854 Py_DECREF(it);
3855 return NULL;
3856 }
3857 Py_INCREF(func);
3858 lz->func = func;
3859 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003860
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003861 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003862}
3863
3864static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003865filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003866{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003867 PyObject_GC_UnTrack(lz);
3868 Py_XDECREF(lz->func);
3869 Py_XDECREF(lz->it);
3870 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003871}
3872
3873static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003874filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003875{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003876 Py_VISIT(lz->it);
3877 Py_VISIT(lz->func);
3878 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003879}
3880
3881static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003882filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003883{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003884 PyObject *item;
3885 PyObject *it = lz->it;
3886 long ok;
3887 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00003888
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003889 iternext = *Py_TYPE(it)->tp_iternext;
3890 for (;;) {
3891 item = iternext(it);
3892 if (item == NULL)
3893 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003894
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003895 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3896 ok = PyObject_IsTrue(item);
3897 } else {
3898 PyObject *good;
Petr Viktorinffd97532020-02-11 17:46:57 +01003899 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003900 if (good == NULL) {
3901 Py_DECREF(item);
3902 return NULL;
3903 }
3904 ok = PyObject_IsTrue(good);
3905 Py_DECREF(good);
3906 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003907 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003908 return item;
3909 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02003910 if (ok < 0)
3911 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003912 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00003913}
3914
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003915static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303916filterfalse_reduce(filterfalseobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003917{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003918 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3919}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003920
3921static PyMethodDef filterfalse_methods[] = {
3922 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3923 reduce_doc},
3924 {NULL, NULL} /* sentinel */
3925};
3926
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003927static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003928 PyVarObject_HEAD_INIT(NULL, 0)
3929 "itertools.filterfalse", /* tp_name */
3930 sizeof(filterfalseobject), /* tp_basicsize */
3931 0, /* tp_itemsize */
3932 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003933 (destructor)filterfalse_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003934 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003935 0, /* tp_getattr */
3936 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003937 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003938 0, /* tp_repr */
3939 0, /* tp_as_number */
3940 0, /* tp_as_sequence */
3941 0, /* tp_as_mapping */
3942 0, /* tp_hash */
3943 0, /* tp_call */
3944 0, /* tp_str */
3945 PyObject_GenericGetAttr, /* tp_getattro */
3946 0, /* tp_setattro */
3947 0, /* tp_as_buffer */
3948 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3949 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003950 itertools_filterfalse__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003951 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003952 0, /* tp_clear */
3953 0, /* tp_richcompare */
3954 0, /* tp_weaklistoffset */
3955 PyObject_SelfIter, /* tp_iter */
3956 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003957 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003958 0, /* tp_members */
3959 0, /* tp_getset */
3960 0, /* tp_base */
3961 0, /* tp_dict */
3962 0, /* tp_descr_get */
3963 0, /* tp_descr_set */
3964 0, /* tp_dictoffset */
3965 0, /* tp_init */
3966 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003967 itertools_filterfalse, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003968 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003969};
3970
3971
3972/* count object ************************************************************/
3973
3974typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003975 PyObject_HEAD
3976 Py_ssize_t cnt;
3977 PyObject *long_cnt;
3978 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00003979} countobject;
3980
Raymond Hettinger30729212009-02-12 06:28:27 +00003981/* Counting logic and invariants:
3982
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003983fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00003984
Serhiy Storchaka483405b2015-02-17 10:14:30 +02003985 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003986 Advances with: cnt += 1
3987 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00003988
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00003989slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00003990
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003991 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3992 All counting is done with python objects (no overflows or underflows).
3993 Advances with: long_cnt += long_step
3994 Step may be zero -- effectively a slow version of repeat(cnt).
3995 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00003996*/
3997
Tal Einatc4bccd32018-09-12 00:49:13 +03003998/*[clinic input]
3999@classmethod
4000itertools.count.__new__
4001 start as long_cnt: object(c_default="NULL") = 0
4002 step as long_step: object(c_default="NULL") = 1
4003Return a count object whose .__next__() method returns consecutive values.
4004
4005Equivalent to:
4006 def count(firstval=0, step=1):
4007 x = firstval
4008 while 1:
4009 yield x
4010 x += step
4011[clinic start generated code]*/
4012
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004013static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03004014itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
4015 PyObject *long_step)
4016/*[clinic end generated code: output=09a9250aebd00b1c input=d7a85eec18bfcd94]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004017{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004018 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004019 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004020 Py_ssize_t cnt = 0;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00004021 long step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004022
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004023 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
4024 (long_step != NULL && !PyNumber_Check(long_step))) {
4025 PyErr_SetString(PyExc_TypeError, "a number is required");
4026 return NULL;
4027 }
Raymond Hettinger30729212009-02-12 06:28:27 +00004028
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004029 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
4030 (long_step == NULL || PyLong_Check(long_step));
4031
4032 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004033 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004034 if (fast_mode) {
4035 assert(PyLong_Check(long_cnt));
4036 cnt = PyLong_AsSsize_t(long_cnt);
4037 if (cnt == -1 && PyErr_Occurred()) {
4038 PyErr_Clear();
4039 fast_mode = 0;
4040 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004041 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004042 } else {
4043 cnt = 0;
Victor Stinner37834132020-10-27 17:12:53 +01004044 long_cnt = _PyLong_GetZero();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004045 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004046 Py_INCREF(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004047
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004048 /* If not specified, step defaults to 1 */
Victor Stinner37834132020-10-27 17:12:53 +01004049 if (long_step == NULL) {
4050 long_step = _PyLong_GetOne();
4051 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004052 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004053
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004054 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004055
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004056 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004057 if (fast_mode) {
4058 assert(PyLong_Check(long_step));
4059 step = PyLong_AsLong(long_step);
4060 if (step != 1) {
4061 fast_mode = 0;
4062 if (step == -1 && PyErr_Occurred())
4063 PyErr_Clear();
4064 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004065 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004066
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004067 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004068 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004069 else
4070 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004071
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004072 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4073 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4074 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004075 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004076
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004077 /* create countobject structure */
4078 lz = (countobject *)type->tp_alloc(type, 0);
4079 if (lz == NULL) {
4080 Py_XDECREF(long_cnt);
Zackery Spytz0523c392019-03-26 00:05:29 -06004081 Py_DECREF(long_step);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004082 return NULL;
4083 }
4084 lz->cnt = cnt;
4085 lz->long_cnt = long_cnt;
4086 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004087
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004088 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004089}
4090
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004091static void
4092count_dealloc(countobject *lz)
4093{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004094 PyObject_GC_UnTrack(lz);
4095 Py_XDECREF(lz->long_cnt);
4096 Py_XDECREF(lz->long_step);
4097 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004098}
4099
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004100static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004101count_traverse(countobject *lz, visitproc visit, void *arg)
4102{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004103 Py_VISIT(lz->long_cnt);
4104 Py_VISIT(lz->long_step);
4105 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004106}
4107
4108static PyObject *
4109count_nextlong(countobject *lz)
4110{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004111 PyObject *long_cnt;
4112 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004113
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004114 long_cnt = lz->long_cnt;
4115 if (long_cnt == NULL) {
4116 /* Switch to slow_mode */
4117 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4118 if (long_cnt == NULL)
4119 return NULL;
4120 }
4121 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004122
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004123 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4124 if (stepped_up == NULL)
4125 return NULL;
4126 lz->long_cnt = stepped_up;
4127 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004128}
4129
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004130static PyObject *
4131count_next(countobject *lz)
4132{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004133 if (lz->cnt == PY_SSIZE_T_MAX)
4134 return count_nextlong(lz);
4135 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004136}
4137
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004138static PyObject *
4139count_repr(countobject *lz)
4140{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004141 if (lz->cnt != PY_SSIZE_T_MAX)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004142 return PyUnicode_FromFormat("%s(%zd)",
4143 _PyType_Name(Py_TYPE(lz)), lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004144
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004145 if (PyLong_Check(lz->long_step)) {
4146 long step = PyLong_AsLong(lz->long_step);
4147 if (step == -1 && PyErr_Occurred()) {
4148 PyErr_Clear();
4149 }
4150 if (step == 1) {
4151 /* Don't display step when it is an integer equal to 1 */
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004152 return PyUnicode_FromFormat("%s(%R)",
4153 _PyType_Name(Py_TYPE(lz)),
4154 lz->long_cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004155 }
4156 }
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004157 return PyUnicode_FromFormat("%s(%R, %R)",
4158 _PyType_Name(Py_TYPE(lz)),
4159 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004160}
4161
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004162static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304163count_reduce(countobject *lz, PyObject *Py_UNUSED(ignored))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004164{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004165 if (lz->cnt == PY_SSIZE_T_MAX)
4166 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4167 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004168}
4169
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004170static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004171 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004172 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004173 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004174};
4175
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004176static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004177 PyVarObject_HEAD_INIT(NULL, 0)
4178 "itertools.count", /* tp_name */
4179 sizeof(countobject), /* tp_basicsize */
4180 0, /* tp_itemsize */
4181 /* methods */
4182 (destructor)count_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004183 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004184 0, /* tp_getattr */
4185 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004186 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004187 (reprfunc)count_repr, /* tp_repr */
4188 0, /* tp_as_number */
4189 0, /* tp_as_sequence */
4190 0, /* tp_as_mapping */
4191 0, /* tp_hash */
4192 0, /* tp_call */
4193 0, /* tp_str */
4194 PyObject_GenericGetAttr, /* tp_getattro */
4195 0, /* tp_setattro */
4196 0, /* tp_as_buffer */
4197 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004198 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03004199 itertools_count__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004200 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004201 0, /* tp_clear */
4202 0, /* tp_richcompare */
4203 0, /* tp_weaklistoffset */
4204 PyObject_SelfIter, /* tp_iter */
4205 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004206 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004207 0, /* tp_members */
4208 0, /* tp_getset */
4209 0, /* tp_base */
4210 0, /* tp_dict */
4211 0, /* tp_descr_get */
4212 0, /* tp_descr_set */
4213 0, /* tp_dictoffset */
4214 0, /* tp_init */
4215 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03004216 itertools_count, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004217 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004218};
4219
4220
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004221/* repeat object ************************************************************/
4222
4223typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004224 PyObject_HEAD
4225 PyObject *element;
4226 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004227} repeatobject;
4228
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004229static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004230
4231static PyObject *
4232repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4233{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004234 repeatobject *ro;
4235 PyObject *element;
Raymond Hettingeraf464502019-11-09 20:28:31 -08004236 Py_ssize_t cnt = -1, n_args;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004237 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004238
Raymond Hettingeraf464502019-11-09 20:28:31 -08004239 n_args = PyTuple_GET_SIZE(args);
4240 if (kwds != NULL)
4241 n_args += PyDict_GET_SIZE(kwds);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004242 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4243 &element, &cnt))
4244 return NULL;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004245 /* Does user supply times argument? */
Raymond Hettingeraf464502019-11-09 20:28:31 -08004246 if (n_args == 2 && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004247 cnt = 0;
4248
4249 ro = (repeatobject *)type->tp_alloc(type, 0);
4250 if (ro == NULL)
4251 return NULL;
4252 Py_INCREF(element);
4253 ro->element = element;
4254 ro->cnt = cnt;
4255 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004256}
4257
4258static void
4259repeat_dealloc(repeatobject *ro)
4260{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004261 PyObject_GC_UnTrack(ro);
4262 Py_XDECREF(ro->element);
4263 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004264}
4265
4266static int
4267repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4268{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004269 Py_VISIT(ro->element);
4270 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004271}
4272
4273static PyObject *
4274repeat_next(repeatobject *ro)
4275{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004276 if (ro->cnt == 0)
4277 return NULL;
4278 if (ro->cnt > 0)
4279 ro->cnt--;
4280 Py_INCREF(ro->element);
4281 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004282}
4283
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004284static PyObject *
4285repeat_repr(repeatobject *ro)
4286{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004287 if (ro->cnt == -1)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004288 return PyUnicode_FromFormat("%s(%R)",
4289 _PyType_Name(Py_TYPE(ro)), ro->element);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004290 else
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004291 return PyUnicode_FromFormat("%s(%R, %zd)",
4292 _PyType_Name(Py_TYPE(ro)), ro->element,
4293 ro->cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004294}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004295
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004296static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304297repeat_len(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004298{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004299 if (ro->cnt == -1) {
4300 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4301 return NULL;
4302 }
4303 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004304}
4305
Armin Rigof5b3e362006-02-11 21:32:43 +00004306PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004307
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004308static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304309repeat_reduce(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004310{
4311 /* unpickle this so that a new repeat iterator is constructed with an
4312 * object, then call __setstate__ on it to set cnt
4313 */
4314 if (ro->cnt >= 0)
4315 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4316 else
4317 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4318}
4319
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004320static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004321 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004322 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004323 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004324};
4325
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004326PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004327"repeat(object [,times]) -> create an iterator which returns the object\n\
4328for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004329endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004330
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004331static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004332 PyVarObject_HEAD_INIT(NULL, 0)
4333 "itertools.repeat", /* tp_name */
4334 sizeof(repeatobject), /* tp_basicsize */
4335 0, /* tp_itemsize */
4336 /* methods */
4337 (destructor)repeat_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004338 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004339 0, /* tp_getattr */
4340 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004341 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004342 (reprfunc)repeat_repr, /* tp_repr */
4343 0, /* tp_as_number */
4344 0, /* tp_as_sequence */
4345 0, /* tp_as_mapping */
4346 0, /* tp_hash */
4347 0, /* tp_call */
4348 0, /* tp_str */
4349 PyObject_GenericGetAttr, /* tp_getattro */
4350 0, /* tp_setattro */
4351 0, /* tp_as_buffer */
4352 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4353 Py_TPFLAGS_BASETYPE, /* tp_flags */
4354 repeat_doc, /* tp_doc */
4355 (traverseproc)repeat_traverse, /* tp_traverse */
4356 0, /* tp_clear */
4357 0, /* tp_richcompare */
4358 0, /* tp_weaklistoffset */
4359 PyObject_SelfIter, /* tp_iter */
4360 (iternextfunc)repeat_next, /* tp_iternext */
4361 repeat_methods, /* tp_methods */
4362 0, /* tp_members */
4363 0, /* tp_getset */
4364 0, /* tp_base */
4365 0, /* tp_dict */
4366 0, /* tp_descr_get */
4367 0, /* tp_descr_set */
4368 0, /* tp_dictoffset */
4369 0, /* tp_init */
4370 0, /* tp_alloc */
4371 repeat_new, /* tp_new */
4372 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004373};
4374
Tal Einatc4bccd32018-09-12 00:49:13 +03004375
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004376/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004377
4378typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004379 PyObject_HEAD
4380 Py_ssize_t tuplesize;
4381 Py_ssize_t numactive;
4382 PyObject *ittuple; /* tuple of iterators */
4383 PyObject *result;
4384 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004385} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004386
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004387static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004388
4389static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004390zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004391{
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004392 _Py_IDENTIFIER(fillvalue);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004393 ziplongestobject *lz;
4394 Py_ssize_t i;
4395 PyObject *ittuple; /* tuple of iterators */
4396 PyObject *result;
4397 PyObject *fillvalue = Py_None;
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004398 Py_ssize_t tuplesize;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004399
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004400 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004401 fillvalue = NULL;
4402 if (PyDict_GET_SIZE(kwds) == 1) {
4403 fillvalue = _PyDict_GetItemIdWithError(kwds, &PyId_fillvalue);
4404 }
4405 if (fillvalue == NULL) {
4406 if (!PyErr_Occurred()) {
4407 PyErr_SetString(PyExc_TypeError,
4408 "zip_longest() got an unexpected keyword argument");
4409 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004410 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004411 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004412 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004413
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004414 /* args must be a tuple */
4415 assert(PyTuple_Check(args));
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004416 tuplesize = PyTuple_GET_SIZE(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004417
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004418 /* obtain iterators */
4419 ittuple = PyTuple_New(tuplesize);
4420 if (ittuple == NULL)
4421 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004422 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004423 PyObject *item = PyTuple_GET_ITEM(args, i);
4424 PyObject *it = PyObject_GetIter(item);
4425 if (it == NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004426 Py_DECREF(ittuple);
4427 return NULL;
4428 }
4429 PyTuple_SET_ITEM(ittuple, i, it);
4430 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004431
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004432 /* create a result holder */
4433 result = PyTuple_New(tuplesize);
4434 if (result == NULL) {
4435 Py_DECREF(ittuple);
4436 return NULL;
4437 }
4438 for (i=0 ; i < tuplesize ; i++) {
4439 Py_INCREF(Py_None);
4440 PyTuple_SET_ITEM(result, i, Py_None);
4441 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004442
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004443 /* create ziplongestobject structure */
4444 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4445 if (lz == NULL) {
4446 Py_DECREF(ittuple);
4447 Py_DECREF(result);
4448 return NULL;
4449 }
4450 lz->ittuple = ittuple;
4451 lz->tuplesize = tuplesize;
4452 lz->numactive = tuplesize;
4453 lz->result = result;
4454 Py_INCREF(fillvalue);
4455 lz->fillvalue = fillvalue;
4456 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004457}
4458
4459static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004460zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004461{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004462 PyObject_GC_UnTrack(lz);
4463 Py_XDECREF(lz->ittuple);
4464 Py_XDECREF(lz->result);
4465 Py_XDECREF(lz->fillvalue);
4466 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004467}
4468
4469static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004470zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004471{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004472 Py_VISIT(lz->ittuple);
4473 Py_VISIT(lz->result);
4474 Py_VISIT(lz->fillvalue);
4475 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004476}
4477
4478static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004479zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004480{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004481 Py_ssize_t i;
4482 Py_ssize_t tuplesize = lz->tuplesize;
4483 PyObject *result = lz->result;
4484 PyObject *it;
4485 PyObject *item;
4486 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004487
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004488 if (tuplesize == 0)
4489 return NULL;
4490 if (lz->numactive == 0)
4491 return NULL;
4492 if (Py_REFCNT(result) == 1) {
4493 Py_INCREF(result);
4494 for (i=0 ; i < tuplesize ; i++) {
4495 it = PyTuple_GET_ITEM(lz->ittuple, i);
4496 if (it == NULL) {
4497 Py_INCREF(lz->fillvalue);
4498 item = lz->fillvalue;
4499 } else {
4500 item = PyIter_Next(it);
4501 if (item == NULL) {
4502 lz->numactive -= 1;
4503 if (lz->numactive == 0 || PyErr_Occurred()) {
4504 lz->numactive = 0;
4505 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004506 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004507 } else {
4508 Py_INCREF(lz->fillvalue);
4509 item = lz->fillvalue;
4510 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4511 Py_DECREF(it);
4512 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004513 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004514 }
4515 olditem = PyTuple_GET_ITEM(result, i);
4516 PyTuple_SET_ITEM(result, i, item);
4517 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004518 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004519 } else {
4520 result = PyTuple_New(tuplesize);
4521 if (result == NULL)
4522 return NULL;
4523 for (i=0 ; i < tuplesize ; i++) {
4524 it = PyTuple_GET_ITEM(lz->ittuple, i);
4525 if (it == NULL) {
4526 Py_INCREF(lz->fillvalue);
4527 item = lz->fillvalue;
4528 } else {
4529 item = PyIter_Next(it);
4530 if (item == NULL) {
4531 lz->numactive -= 1;
4532 if (lz->numactive == 0 || PyErr_Occurred()) {
4533 lz->numactive = 0;
4534 Py_DECREF(result);
4535 return NULL;
4536 } else {
4537 Py_INCREF(lz->fillvalue);
4538 item = lz->fillvalue;
4539 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4540 Py_DECREF(it);
4541 }
4542 }
4543 }
4544 PyTuple_SET_ITEM(result, i, item);
4545 }
4546 }
4547 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004548}
4549
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004550static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304551zip_longest_reduce(ziplongestobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004552{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004553
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004554 /* Create a new tuple with empty sequences where appropriate to pickle.
4555 * Then use setstate to set the fillvalue
4556 */
4557 int i;
4558 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004559
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004560 if (args == NULL)
4561 return NULL;
4562 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4563 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4564 if (elem == NULL) {
4565 elem = PyTuple_New(0);
4566 if (elem == NULL) {
4567 Py_DECREF(args);
4568 return NULL;
4569 }
4570 } else
4571 Py_INCREF(elem);
4572 PyTuple_SET_ITEM(args, i, elem);
4573 }
4574 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4575}
4576
4577static PyObject *
4578zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4579{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004580 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004581 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004582 Py_RETURN_NONE;
4583}
4584
4585static PyMethodDef zip_longest_methods[] = {
4586 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4587 reduce_doc},
4588 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4589 setstate_doc},
4590 {NULL, NULL} /* sentinel */
4591};
4592
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004593PyDoc_STRVAR(zip_longest_doc,
4594"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004595\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004596Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004597the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004598method continues until the longest iterable in the argument sequence\n\
4599is exhausted and then it raises StopIteration. When the shorter iterables\n\
4600are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4601defaults to None or can be specified by a keyword argument.\n\
4602");
4603
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004604static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004605 PyVarObject_HEAD_INIT(NULL, 0)
4606 "itertools.zip_longest", /* tp_name */
4607 sizeof(ziplongestobject), /* tp_basicsize */
4608 0, /* tp_itemsize */
4609 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004610 (destructor)zip_longest_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004611 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004612 0, /* tp_getattr */
4613 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004614 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004615 0, /* tp_repr */
4616 0, /* tp_as_number */
4617 0, /* tp_as_sequence */
4618 0, /* tp_as_mapping */
4619 0, /* tp_hash */
4620 0, /* tp_call */
4621 0, /* tp_str */
4622 PyObject_GenericGetAttr, /* tp_getattro */
4623 0, /* tp_setattro */
4624 0, /* tp_as_buffer */
4625 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4626 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004627 zip_longest_doc, /* tp_doc */
4628 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004629 0, /* tp_clear */
4630 0, /* tp_richcompare */
4631 0, /* tp_weaklistoffset */
4632 PyObject_SelfIter, /* tp_iter */
4633 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004634 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004635 0, /* tp_members */
4636 0, /* tp_getset */
4637 0, /* tp_base */
4638 0, /* tp_dict */
4639 0, /* tp_descr_get */
4640 0, /* tp_descr_set */
4641 0, /* tp_dictoffset */
4642 0, /* tp_init */
4643 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004644 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004645 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004646};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004647
Tal Einatc4bccd32018-09-12 00:49:13 +03004648
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004649/* module level code ********************************************************/
4650
4651PyDoc_STRVAR(module_doc,
4652"Functional tools for creating and using iterators.\n\
4653\n\
4654Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004655count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004656cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004657repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004658\n\
4659Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004660accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
oldkaa0735f2018-02-02 16:52:55 +08004661chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
4662chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004663compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4664dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4665groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004666filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004667islice(seq, [start,] stop [, step]) --> elements from\n\
4668 seq[start:stop:step]\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004669starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004670tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004671takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
oldkaa0735f2018-02-02 16:52:55 +08004672zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004673\n\
4674Combinatoric generators:\n\
4675product(p, q, ... [repeat=1]) --> cartesian product\n\
4676permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004677combinations(p, r)\n\
4678combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004679");
4680
Dong-hee Na514c4692020-03-18 02:46:24 +09004681static int
4682itertoolsmodule_exec(PyObject *m)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004683{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004684 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004685 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004686 &combinations_type,
4687 &cwr_type,
4688 &cycle_type,
4689 &dropwhile_type,
4690 &takewhile_type,
4691 &islice_type,
4692 &starmap_type,
4693 &chain_type,
4694 &compress_type,
4695 &filterfalse_type,
4696 &count_type,
4697 &ziplongest_type,
4698 &permutations_type,
4699 &product_type,
4700 &repeat_type,
4701 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004702 &_grouper_type,
4703 &tee_type,
Dong-hee Na05e4a292020-03-23 01:17:34 +09004704 &teedataobject_type
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004705 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004706
Victor Stinnerd2ec81a2020-02-07 09:17:07 +01004707 Py_SET_TYPE(&teedataobject_type, &PyType_Type);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004708
Dong-hee Na05e4a292020-03-23 01:17:34 +09004709 for (size_t i = 0; i < Py_ARRAY_LENGTH(typelist); i++) {
4710 if (PyModule_AddType(m, typelist[i]) < 0) {
Dong-hee Na514c4692020-03-18 02:46:24 +09004711 return -1;
4712 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004713 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004714
Dong-hee Na514c4692020-03-18 02:46:24 +09004715 return 0;
4716}
4717
4718static struct PyModuleDef_Slot itertoolsmodule_slots[] = {
4719 {Py_mod_exec, itertoolsmodule_exec},
4720 {0, NULL}
4721};
4722
4723static PyMethodDef module_methods[] = {
4724 ITERTOOLS_TEE_METHODDEF
4725 {NULL, NULL} /* sentinel */
4726};
4727
4728
4729static struct PyModuleDef itertoolsmodule = {
4730 PyModuleDef_HEAD_INIT,
4731 "itertools",
4732 module_doc,
4733 0,
4734 module_methods,
4735 itertoolsmodule_slots,
4736 NULL,
4737 NULL,
4738 NULL
4739};
4740
4741PyMODINIT_FUNC
4742PyInit_itertools(void)
4743{
4744 return PyModuleDef_Init(&itertoolsmodule);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004745}