blob: 7144856c352f81f6d5f60488d6de2c240b590cfc [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001
Raymond Hettingercc061d02020-11-30 20:42:54 -08002
Raymond Hettingerca3788c2015-08-16 14:49:24 -07003#define PY_SSIZE_T_CLEAN
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004#include "Python.h"
Victor Stinner37834132020-10-27 17:12:53 +01005#include "pycore_long.h" // _PyLong_GetZero()
Victor Stinner384621c2020-06-22 17:27:35 +02006#include "pycore_tuple.h" // _PyTuple_ITEMS()
Victor Stinner4a21e572020-04-15 02:35:41 +02007#include <stddef.h> // offsetof()
Raymond Hettinger96ef8112003-02-01 00:10:11 +00008
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00009/* Itertools module written and maintained
Raymond Hettinger96ef8112003-02-01 00:10:11 +000010 by Raymond D. Hettinger <python@rcn.com>
Raymond Hettinger96ef8112003-02-01 00:10:11 +000011*/
12
Tal Einat3286ce42018-09-10 21:33:08 +030013/*[clinic input]
14module itertools
15class itertools.groupby "groupbyobject *" "&groupby_type"
16class itertools._grouper "_grouperobject *" "&_grouper_type"
Tal Einatc4bccd32018-09-12 00:49:13 +030017class itertools.teedataobject "teedataobject *" "&teedataobject_type"
18class itertools._tee "teeobject *" "&tee_type"
19class itertools.cycle "cycleobject *" "&cycle_type"
20class itertools.dropwhile "dropwhileobject *" "&dropwhile_type"
21class itertools.takewhile "takewhileobject *" "&takewhile_type"
22class itertools.starmap "starmapobject *" "&starmap_type"
23class itertools.chain "chainobject *" "&chain_type"
24class itertools.combinations "combinationsobject *" "&combinations_type"
25class itertools.combinations_with_replacement "cwr_object *" "&cwr_type"
26class itertools.permutations "permutationsobject *" "&permutations_type"
27class itertools.accumulate "accumulateobject *" "&accumulate_type"
28class itertools.compress "compressobject *" "&compress_type"
29class itertools.filterfalse "filterfalseobject *" "&filterfalse_type"
30class itertools.count "countobject *" "&count_type"
Raymond Hettingercc061d02020-11-30 20:42:54 -080031class itertools.pairwise "pairwiseobject *" "&pairwise_type"
Tal Einat3286ce42018-09-10 21:33:08 +030032[clinic start generated code]*/
Raymond Hettingercc061d02020-11-30 20:42:54 -080033/*[clinic end generated code: output=da39a3ee5e6b4b0d input=6498ed21fbe1bf94]*/
Tal Einat3286ce42018-09-10 21:33:08 +030034
35static PyTypeObject groupby_type;
36static PyTypeObject _grouper_type;
Tal Einatc4bccd32018-09-12 00:49:13 +030037static PyTypeObject teedataobject_type;
38static PyTypeObject tee_type;
39static PyTypeObject cycle_type;
40static PyTypeObject dropwhile_type;
41static PyTypeObject takewhile_type;
42static PyTypeObject starmap_type;
43static PyTypeObject combinations_type;
44static PyTypeObject cwr_type;
45static PyTypeObject permutations_type;
46static PyTypeObject accumulate_type;
47static PyTypeObject compress_type;
48static PyTypeObject filterfalse_type;
49static PyTypeObject count_type;
Raymond Hettingercc061d02020-11-30 20:42:54 -080050static PyTypeObject pairwise_type;
Tal Einatc4bccd32018-09-12 00:49:13 +030051
Tal Einat3286ce42018-09-10 21:33:08 +030052#include "clinic/itertoolsmodule.c.h"
53
Raymond Hettingercc061d02020-11-30 20:42:54 -080054/* pairwise object ***********************************************************/
55
56typedef struct {
57 PyObject_HEAD
58 PyObject *it;
59 PyObject *old;
60} pairwiseobject;
61
62/*[clinic input]
63@classmethod
64itertools.pairwise.__new__ as pairwise_new
65 iterable: object
66 /
67Return an iterator of overlapping pairs taken from the input iterator.
68
69 s -> (s0,s1), (s1,s2), (s2, s3), ...
70
71[clinic start generated code]*/
72
73static PyObject *
74pairwise_new_impl(PyTypeObject *type, PyObject *iterable)
75/*[clinic end generated code: output=9f0267062d384456 input=6e7c3cddb431a8d6]*/
76{
77 PyObject *it;
78 pairwiseobject *po;
79
80 it = PyObject_GetIter(iterable);
81 if (it == NULL) {
82 return NULL;
83 }
84 po = (pairwiseobject *)type->tp_alloc(type, 0);
85 if (po == NULL) {
86 Py_DECREF(it);
87 return NULL;
88 }
89 po->it = it;
90 po->old = NULL;
91 return (PyObject *)po;
92}
93
94static void
95pairwise_dealloc(pairwiseobject *po)
96{
97 PyObject_GC_UnTrack(po);
98 Py_XDECREF(po->it);
99 Py_XDECREF(po->old);
100 Py_TYPE(po)->tp_free(po);
101}
102
103static int
104pairwise_traverse(pairwiseobject *po, visitproc visit, void *arg)
105{
106 Py_VISIT(po->it);
107 Py_VISIT(po->old);
108 return 0;
109}
110
111static PyObject *
112pairwise_next(pairwiseobject *po)
113{
114 PyObject *it = po->it;
115 PyObject *old = po->old;
116 PyObject *new, *result;
117
118 if (it == NULL) {
119 return NULL;
120 }
121 if (old == NULL) {
122 po->old = old = (*Py_TYPE(it)->tp_iternext)(it);
123 if (old == NULL) {
124 Py_CLEAR(po->it);
125 return NULL;
126 }
127 }
128 new = (*Py_TYPE(it)->tp_iternext)(it);
129 if (new == NULL) {
130 Py_CLEAR(po->it);
131 Py_CLEAR(po->old);
132 return NULL;
133 }
134 /* Future optimization: Reuse the result tuple as we do in enumerate() */
135 result = PyTuple_Pack(2, old, new);
136 Py_SETREF(po->old, new);
137 return result;
138}
139
140static PyTypeObject pairwise_type = {
141 PyVarObject_HEAD_INIT(&PyType_Type, 0)
142 "itertools.pairwise", /* tp_name */
143 sizeof(pairwiseobject), /* tp_basicsize */
144 0, /* tp_itemsize */
145 /* methods */
146 (destructor)pairwise_dealloc, /* tp_dealloc */
147 0, /* tp_vectorcall_offset */
148 0, /* tp_getattr */
149 0, /* tp_setattr */
150 0, /* tp_as_async */
151 0, /* tp_repr */
152 0, /* tp_as_number */
153 0, /* tp_as_sequence */
154 0, /* tp_as_mapping */
155 0, /* tp_hash */
156 0, /* tp_call */
157 0, /* tp_str */
158 PyObject_GenericGetAttr, /* tp_getattro */
159 0, /* tp_setattro */
160 0, /* tp_as_buffer */
161 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
162 Py_TPFLAGS_BASETYPE, /* tp_flags */
163 pairwise_new__doc__, /* tp_doc */
164 (traverseproc)pairwise_traverse, /* tp_traverse */
165 0, /* tp_clear */
166 0, /* tp_richcompare */
167 0, /* tp_weaklistoffset */
168 PyObject_SelfIter, /* tp_iter */
169 (iternextfunc)pairwise_next, /* tp_iternext */
170 0, /* tp_methods */
171 0, /* tp_members */
172 0, /* tp_getset */
173 0, /* tp_base */
174 0, /* tp_dict */
175 0, /* tp_descr_get */
176 0, /* tp_descr_set */
177 0, /* tp_dictoffset */
178 0, /* tp_init */
179 PyType_GenericAlloc, /* tp_alloc */
180 pairwise_new, /* tp_new */
181 PyObject_GC_Del, /* tp_free */
182};
183
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000184
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700185/* groupby object ************************************************************/
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000186
187typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000188 PyObject_HEAD
189 PyObject *it;
190 PyObject *keyfunc;
191 PyObject *tgtkey;
192 PyObject *currkey;
193 PyObject *currvalue;
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300194 const void *currgrouper; /* borrowed reference */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000195} groupbyobject;
196
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000197static PyObject *_grouper_create(groupbyobject *, PyObject *);
198
Tal Einat3286ce42018-09-10 21:33:08 +0300199/*[clinic input]
200@classmethod
201itertools.groupby.__new__
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000202
Tal Einat3286ce42018-09-10 21:33:08 +0300203 iterable as it: object
204 Elements to divide into groups according to the key function.
205 key as keyfunc: object = None
206 A function for computing the group category for each element.
207 If the key function is not specified or is None, the element itself
208 is used for grouping.
209
210make an iterator that returns consecutive keys and groups from the iterable
211[clinic start generated code]*/
212
213static PyObject *
214itertools_groupby_impl(PyTypeObject *type, PyObject *it, PyObject *keyfunc)
215/*[clinic end generated code: output=cbb1ae3a90fd4141 input=6b3d123e87ff65a1]*/
216{
217 groupbyobject *gbo;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000218
219 gbo = (groupbyobject *)type->tp_alloc(type, 0);
220 if (gbo == NULL)
221 return NULL;
222 gbo->tgtkey = NULL;
223 gbo->currkey = NULL;
224 gbo->currvalue = NULL;
225 gbo->keyfunc = keyfunc;
226 Py_INCREF(keyfunc);
227 gbo->it = PyObject_GetIter(it);
228 if (gbo->it == NULL) {
229 Py_DECREF(gbo);
230 return NULL;
231 }
232 return (PyObject *)gbo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000233}
234
235static void
236groupby_dealloc(groupbyobject *gbo)
237{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000238 PyObject_GC_UnTrack(gbo);
239 Py_XDECREF(gbo->it);
240 Py_XDECREF(gbo->keyfunc);
241 Py_XDECREF(gbo->tgtkey);
242 Py_XDECREF(gbo->currkey);
243 Py_XDECREF(gbo->currvalue);
244 Py_TYPE(gbo)->tp_free(gbo);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000245}
246
247static int
248groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
249{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000250 Py_VISIT(gbo->it);
251 Py_VISIT(gbo->keyfunc);
252 Py_VISIT(gbo->tgtkey);
253 Py_VISIT(gbo->currkey);
254 Py_VISIT(gbo->currvalue);
255 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000256}
257
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300258Py_LOCAL_INLINE(int)
259groupby_step(groupbyobject *gbo)
260{
261 PyObject *newvalue, *newkey, *oldvalue;
262
263 newvalue = PyIter_Next(gbo->it);
264 if (newvalue == NULL)
265 return -1;
266
267 if (gbo->keyfunc == Py_None) {
268 newkey = newvalue;
269 Py_INCREF(newvalue);
270 } else {
Petr Viktorinffd97532020-02-11 17:46:57 +0100271 newkey = PyObject_CallOneArg(gbo->keyfunc, newvalue);
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300272 if (newkey == NULL) {
273 Py_DECREF(newvalue);
274 return -1;
275 }
276 }
277
278 oldvalue = gbo->currvalue;
279 gbo->currvalue = newvalue;
280 Py_XSETREF(gbo->currkey, newkey);
281 Py_XDECREF(oldvalue);
282 return 0;
283}
284
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000285static PyObject *
286groupby_next(groupbyobject *gbo)
287{
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300288 PyObject *r, *grouper;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000289
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300290 gbo->currgrouper = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000291 /* skip to next iteration group */
292 for (;;) {
293 if (gbo->currkey == NULL)
294 /* pass */;
295 else if (gbo->tgtkey == NULL)
296 break;
297 else {
298 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000299
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700300 rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000301 if (rcmp == -1)
302 return NULL;
303 else if (rcmp == 0)
304 break;
305 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000306
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300307 if (groupby_step(gbo) < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000308 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000309 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000310 Py_INCREF(gbo->currkey);
Serhiy Storchakaec397562016-04-06 09:50:03 +0300311 Py_XSETREF(gbo->tgtkey, gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000312
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000313 grouper = _grouper_create(gbo, gbo->tgtkey);
314 if (grouper == NULL)
315 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000316
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000317 r = PyTuple_Pack(2, gbo->currkey, grouper);
318 Py_DECREF(grouper);
319 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000320}
321
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000322static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530323groupby_reduce(groupbyobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000324{
325 /* reduce as a 'new' call with an optional 'setstate' if groupby
326 * has started
327 */
328 PyObject *value;
329 if (lz->tgtkey && lz->currkey && lz->currvalue)
330 value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
331 lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
332 else
333 value = Py_BuildValue("O(OO)", Py_TYPE(lz),
334 lz->it, lz->keyfunc);
335
336 return value;
337}
338
339PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
340
341static PyObject *
342groupby_setstate(groupbyobject *lz, PyObject *state)
343{
344 PyObject *currkey, *currvalue, *tgtkey;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300345 if (!PyTuple_Check(state)) {
346 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000347 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300348 }
349 if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey)) {
350 return NULL;
351 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200352 Py_INCREF(currkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300353 Py_XSETREF(lz->currkey, currkey);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200354 Py_INCREF(currvalue);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300355 Py_XSETREF(lz->currvalue, currvalue);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200356 Py_INCREF(tgtkey);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300357 Py_XSETREF(lz->tgtkey, tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000358 Py_RETURN_NONE;
359}
360
361PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
362
363static PyMethodDef groupby_methods[] = {
364 {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
365 reduce_doc},
366 {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
367 setstate_doc},
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700368 {NULL, NULL} /* sentinel */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000369};
370
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000371static PyTypeObject groupby_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000372 PyVarObject_HEAD_INIT(NULL, 0)
373 "itertools.groupby", /* tp_name */
374 sizeof(groupbyobject), /* tp_basicsize */
375 0, /* tp_itemsize */
376 /* methods */
377 (destructor)groupby_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200378 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000379 0, /* tp_getattr */
380 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200381 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000382 0, /* tp_repr */
383 0, /* tp_as_number */
384 0, /* tp_as_sequence */
385 0, /* tp_as_mapping */
386 0, /* tp_hash */
387 0, /* tp_call */
388 0, /* tp_str */
389 PyObject_GenericGetAttr, /* tp_getattro */
390 0, /* tp_setattro */
391 0, /* tp_as_buffer */
392 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
393 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einat3286ce42018-09-10 21:33:08 +0300394 itertools_groupby__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000395 (traverseproc)groupby_traverse, /* tp_traverse */
396 0, /* tp_clear */
397 0, /* tp_richcompare */
398 0, /* tp_weaklistoffset */
399 PyObject_SelfIter, /* tp_iter */
400 (iternextfunc)groupby_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000401 groupby_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000402 0, /* tp_members */
403 0, /* tp_getset */
404 0, /* tp_base */
405 0, /* tp_dict */
406 0, /* tp_descr_get */
407 0, /* tp_descr_set */
408 0, /* tp_dictoffset */
409 0, /* tp_init */
410 0, /* tp_alloc */
Tal Einat3286ce42018-09-10 21:33:08 +0300411 itertools_groupby, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000412 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000413};
414
415
416/* _grouper object (internal) ************************************************/
417
418typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000419 PyObject_HEAD
420 PyObject *parent;
421 PyObject *tgtkey;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000422} _grouperobject;
423
Tal Einat3286ce42018-09-10 21:33:08 +0300424/*[clinic input]
425@classmethod
426itertools._grouper.__new__
427
428 parent: object(subclass_of='&groupby_type')
429 tgtkey: object
430 /
431[clinic start generated code]*/
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000432
433static PyObject *
Tal Einat3286ce42018-09-10 21:33:08 +0300434itertools__grouper_impl(PyTypeObject *type, PyObject *parent,
435 PyObject *tgtkey)
436/*[clinic end generated code: output=462efb1cdebb5914 input=dc180d7771fc8c59]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000437{
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000438 return _grouper_create((groupbyobject*) parent, tgtkey);
439}
440
441static PyObject *
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000442_grouper_create(groupbyobject *parent, PyObject *tgtkey)
443{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000444 _grouperobject *igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000445
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000446 igo = PyObject_GC_New(_grouperobject, &_grouper_type);
447 if (igo == NULL)
448 return NULL;
449 igo->parent = (PyObject *)parent;
450 Py_INCREF(parent);
451 igo->tgtkey = tgtkey;
452 Py_INCREF(tgtkey);
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300453 parent->currgrouper = igo; /* borrowed reference */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000454
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000455 PyObject_GC_Track(igo);
456 return (PyObject *)igo;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000457}
458
459static void
460_grouper_dealloc(_grouperobject *igo)
461{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000462 PyObject_GC_UnTrack(igo);
463 Py_DECREF(igo->parent);
464 Py_DECREF(igo->tgtkey);
465 PyObject_GC_Del(igo);
Christian Heimesdd15f6c2008-03-16 00:07:10 +0000466}
467
468static int
469_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
470{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000471 Py_VISIT(igo->parent);
472 Py_VISIT(igo->tgtkey);
473 return 0;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000474}
475
476static PyObject *
477_grouper_next(_grouperobject *igo)
478{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000479 groupbyobject *gbo = (groupbyobject *)igo->parent;
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300480 PyObject *r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000481 int rcmp;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000482
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300483 if (gbo->currgrouper != igo)
484 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000485 if (gbo->currvalue == NULL) {
Serhiy Storchakac740e4f2017-09-26 21:47:56 +0300486 if (groupby_step(gbo) < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000487 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000488 }
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000489
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000490 assert(gbo->currkey != NULL);
491 rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
492 if (rcmp <= 0)
493 /* got any error or current group is end */
494 return NULL;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000495
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000496 r = gbo->currvalue;
497 gbo->currvalue = NULL;
498 Py_CLEAR(gbo->currkey);
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000499
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000500 return r;
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000501}
502
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000503static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530504_grouper_reduce(_grouperobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000505{
Serhiy Storchakabb86bf42018-12-11 08:28:18 +0200506 _Py_IDENTIFIER(iter);
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300507 if (((groupbyobject *)lz->parent)->currgrouper != lz) {
Serhiy Storchakabb86bf42018-12-11 08:28:18 +0200508 return Py_BuildValue("N(())", _PyEval_GetBuiltinId(&PyId_iter));
Serhiy Storchakac247caf2017-09-24 13:36:11 +0300509 }
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700510 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000511}
512
513static PyMethodDef _grouper_methods[] = {
514 {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
515 reduce_doc},
516 {NULL, NULL} /* sentinel */
517};
518
519
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000520static PyTypeObject _grouper_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000521 PyVarObject_HEAD_INIT(NULL, 0)
522 "itertools._grouper", /* tp_name */
523 sizeof(_grouperobject), /* tp_basicsize */
524 0, /* tp_itemsize */
525 /* methods */
526 (destructor)_grouper_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200527 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000528 0, /* tp_getattr */
529 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200530 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000531 0, /* tp_repr */
532 0, /* tp_as_number */
533 0, /* tp_as_sequence */
534 0, /* tp_as_mapping */
535 0, /* tp_hash */
536 0, /* tp_call */
537 0, /* tp_str */
538 PyObject_GenericGetAttr, /* tp_getattro */
539 0, /* tp_setattro */
540 0, /* tp_as_buffer */
541 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
542 0, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700543 (traverseproc)_grouper_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000544 0, /* tp_clear */
545 0, /* tp_richcompare */
546 0, /* tp_weaklistoffset */
547 PyObject_SelfIter, /* tp_iter */
548 (iternextfunc)_grouper_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000549 _grouper_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000550 0, /* tp_members */
551 0, /* tp_getset */
552 0, /* tp_base */
553 0, /* tp_dict */
554 0, /* tp_descr_get */
555 0, /* tp_descr_set */
556 0, /* tp_dictoffset */
557 0, /* tp_init */
558 0, /* tp_alloc */
Tal Einat3286ce42018-09-10 21:33:08 +0300559 itertools__grouper, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000560 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000561};
562
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000563
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700564/* tee object and with supporting function and objects ***********************/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000565
Raymond Hettingerad983e72003-11-12 14:32:26 +0000566/* The teedataobject pre-allocates space for LINKCELLS number of objects.
567 To help the object fit neatly inside cache lines (space for 16 to 32
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000568 pointers), the value should be a multiple of 16 minus space for
Raymond Hettingerad983e72003-11-12 14:32:26 +0000569 the other structure members including PyHEAD overhead. The larger the
570 value, the less memory overhead per object and the less time spent
571 allocating/deallocating new links. The smaller the number, the less
572 wasted space and the more rapid freeing of older data.
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000573*/
Raymond Hettingerad983e72003-11-12 14:32:26 +0000574#define LINKCELLS 57
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000575
576typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000577 PyObject_HEAD
578 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700579 int numread; /* 0 <= numread <= LINKCELLS */
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300580 int running;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000581 PyObject *nextlink;
582 PyObject *(values[LINKCELLS]);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000583} teedataobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000584
585typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000586 PyObject_HEAD
587 teedataobject *dataobj;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700588 int index; /* 0 <= index <= LINKCELLS */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000589 PyObject *weakreflist;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000590} teeobject;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000591
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000592static PyObject *
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000593teedataobject_newinternal(PyObject *it)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000594{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000595 teedataobject *tdo;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000596
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000597 tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
598 if (tdo == NULL)
599 return NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000600
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300601 tdo->running = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000602 tdo->numread = 0;
603 tdo->nextlink = NULL;
604 Py_INCREF(it);
605 tdo->it = it;
606 PyObject_GC_Track(tdo);
607 return (PyObject *)tdo;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000608}
609
610static PyObject *
611teedataobject_jumplink(teedataobject *tdo)
612{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000613 if (tdo->nextlink == NULL)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000614 tdo->nextlink = teedataobject_newinternal(tdo->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000615 Py_XINCREF(tdo->nextlink);
616 return tdo->nextlink;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000617}
618
619static PyObject *
620teedataobject_getitem(teedataobject *tdo, int i)
621{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000622 PyObject *value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000623
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000624 assert(i < LINKCELLS);
625 if (i < tdo->numread)
626 value = tdo->values[i];
627 else {
628 /* this is the lead iterator, so fetch more data */
629 assert(i == tdo->numread);
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300630 if (tdo->running) {
631 PyErr_SetString(PyExc_RuntimeError,
632 "cannot re-enter the tee iterator");
633 return NULL;
634 }
635 tdo->running = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000636 value = PyIter_Next(tdo->it);
Serhiy Storchaka526a0142019-09-09 11:47:14 +0300637 tdo->running = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000638 if (value == NULL)
639 return NULL;
640 tdo->numread++;
641 tdo->values[i] = value;
642 }
643 Py_INCREF(value);
644 return value;
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000645}
646
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000647static int
648teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
649{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000650 int i;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700651
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000652 Py_VISIT(tdo->it);
653 for (i = 0; i < tdo->numread; i++)
654 Py_VISIT(tdo->values[i]);
655 Py_VISIT(tdo->nextlink);
656 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000657}
658
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200659static void
660teedataobject_safe_decref(PyObject *obj)
661{
Andy Lesterdffe4c02020-03-04 07:15:20 -0600662 while (obj && Py_IS_TYPE(obj, &teedataobject_type) &&
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200663 Py_REFCNT(obj) == 1) {
664 PyObject *nextlink = ((teedataobject *)obj)->nextlink;
665 ((teedataobject *)obj)->nextlink = NULL;
666 Py_DECREF(obj);
667 obj = nextlink;
668 }
669 Py_XDECREF(obj);
670}
671
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000672static int
673teedataobject_clear(teedataobject *tdo)
674{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000675 int i;
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200676 PyObject *tmp;
677
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000678 Py_CLEAR(tdo->it);
679 for (i=0 ; i<tdo->numread ; i++)
680 Py_CLEAR(tdo->values[i]);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200681 tmp = tdo->nextlink;
682 tdo->nextlink = NULL;
683 teedataobject_safe_decref(tmp);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000684 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000685}
686
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000687static void
Raymond Hettingerad983e72003-11-12 14:32:26 +0000688teedataobject_dealloc(teedataobject *tdo)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000689{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000690 PyObject_GC_UnTrack(tdo);
691 teedataobject_clear(tdo);
692 PyObject_GC_Del(tdo);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000693}
694
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000695static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530696teedataobject_reduce(teedataobject *tdo, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000697{
698 int i;
699 /* create a temporary list of already iterated values */
700 PyObject *values = PyList_New(tdo->numread);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -0700701
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000702 if (!values)
703 return NULL;
704 for (i=0 ; i<tdo->numread ; i++) {
705 Py_INCREF(tdo->values[i]);
706 PyList_SET_ITEM(values, i, tdo->values[i]);
707 }
708 return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
709 values,
710 tdo->nextlink ? tdo->nextlink : Py_None);
711}
712
Tal Einatc4bccd32018-09-12 00:49:13 +0300713/*[clinic input]
714@classmethod
715itertools.teedataobject.__new__
716 iterable as it: object
717 values: object(subclass_of='&PyList_Type')
718 next: object
719 /
720Data container common to multiple tee objects.
721[clinic start generated code]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000722
723static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +0300724itertools_teedataobject_impl(PyTypeObject *type, PyObject *it,
725 PyObject *values, PyObject *next)
726/*[clinic end generated code: output=3343ceb07e08df5e input=be60f2fabd2b72ba]*/
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000727{
728 teedataobject *tdo;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000729 Py_ssize_t i, len;
730
731 assert(type == &teedataobject_type);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000732
733 tdo = (teedataobject *)teedataobject_newinternal(it);
734 if (!tdo)
735 return NULL;
736
737 len = PyList_GET_SIZE(values);
738 if (len > LINKCELLS)
739 goto err;
740 for (i=0; i<len; i++) {
741 tdo->values[i] = PyList_GET_ITEM(values, i);
742 Py_INCREF(tdo->values[i]);
743 }
Martin v. Löwis33cac852012-05-15 14:34:58 +0200744 /* len <= LINKCELLS < INT_MAX */
745 tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000746
747 if (len == LINKCELLS) {
748 if (next != Py_None) {
Andy Lester55728702020-03-06 16:53:17 -0600749 if (!Py_IS_TYPE(next, &teedataobject_type))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000750 goto err;
751 assert(tdo->nextlink == NULL);
752 Py_INCREF(next);
753 tdo->nextlink = next;
754 }
755 } else {
756 if (next != Py_None)
757 goto err; /* shouldn't have a next if we are not full */
758 }
759 return (PyObject*)tdo;
760
761err:
762 Py_XDECREF(tdo);
763 PyErr_SetString(PyExc_ValueError, "Invalid arguments");
764 return NULL;
765}
766
767static PyMethodDef teedataobject_methods[] = {
768 {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
769 reduce_doc},
770 {NULL, NULL} /* sentinel */
771};
772
Raymond Hettingerad983e72003-11-12 14:32:26 +0000773static PyTypeObject teedataobject_type = {
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700774 PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000775 "itertools._tee_dataobject", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000776 sizeof(teedataobject), /* tp_basicsize */
777 0, /* tp_itemsize */
778 /* methods */
779 (destructor)teedataobject_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200780 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000781 0, /* tp_getattr */
782 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200783 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000784 0, /* tp_repr */
785 0, /* tp_as_number */
786 0, /* tp_as_sequence */
787 0, /* tp_as_mapping */
788 0, /* tp_hash */
789 0, /* tp_call */
790 0, /* tp_str */
791 PyObject_GenericGetAttr, /* tp_getattro */
792 0, /* tp_setattro */
793 0, /* tp_as_buffer */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700794 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300795 itertools_teedataobject__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000796 (traverseproc)teedataobject_traverse, /* tp_traverse */
797 (inquiry)teedataobject_clear, /* tp_clear */
798 0, /* tp_richcompare */
799 0, /* tp_weaklistoffset */
800 0, /* tp_iter */
801 0, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000802 teedataobject_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000803 0, /* tp_members */
804 0, /* tp_getset */
805 0, /* tp_base */
806 0, /* tp_dict */
807 0, /* tp_descr_get */
808 0, /* tp_descr_set */
809 0, /* tp_dictoffset */
810 0, /* tp_init */
811 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +0300812 itertools_teedataobject, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000813 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000814};
815
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000816
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000817static PyObject *
Raymond Hettingerad983e72003-11-12 14:32:26 +0000818tee_next(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000819{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000820 PyObject *value, *link;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000821
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000822 if (to->index >= LINKCELLS) {
823 link = teedataobject_jumplink(to->dataobj);
Serhiy Storchakaa3e91282013-01-25 13:19:31 +0200824 if (link == NULL)
825 return NULL;
Serhiy Storchaka57a01d32016-04-10 18:05:40 +0300826 Py_SETREF(to->dataobj, (teedataobject *)link);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000827 to->index = 0;
828 }
829 value = teedataobject_getitem(to->dataobj, to->index);
830 if (value == NULL)
831 return NULL;
832 to->index++;
833 return value;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000834}
835
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000836static int
837tee_traverse(teeobject *to, visitproc visit, void *arg)
838{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000839 Py_VISIT((PyObject *)to->dataobj);
840 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000841}
842
Raymond Hettingerad983e72003-11-12 14:32:26 +0000843static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530844tee_copy(teeobject *to, PyObject *Py_UNUSED(ignored))
Raymond Hettingerad983e72003-11-12 14:32:26 +0000845{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000846 teeobject *newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000847
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000848 newto = PyObject_GC_New(teeobject, &tee_type);
849 if (newto == NULL)
850 return NULL;
851 Py_INCREF(to->dataobj);
852 newto->dataobj = to->dataobj;
853 newto->index = to->index;
854 newto->weakreflist = NULL;
855 PyObject_GC_Track(newto);
856 return (PyObject *)newto;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000857}
858
859PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
860
861static PyObject *
862tee_fromiterable(PyObject *iterable)
863{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000864 teeobject *to;
865 PyObject *it = NULL;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000866
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000867 it = PyObject_GetIter(iterable);
868 if (it == NULL)
869 return NULL;
870 if (PyObject_TypeCheck(it, &tee_type)) {
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530871 to = (teeobject *)tee_copy((teeobject *)it, NULL);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000872 goto done;
873 }
Raymond Hettingerad983e72003-11-12 14:32:26 +0000874
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000875 to = PyObject_GC_New(teeobject, &tee_type);
876 if (to == NULL)
877 goto done;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000878 to->dataobj = (teedataobject *)teedataobject_newinternal(it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000879 if (!to->dataobj) {
880 PyObject_GC_Del(to);
881 to = NULL;
882 goto done;
883 }
Thomas Wouters0e3f5912006-08-11 14:57:12 +0000884
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000885 to->index = 0;
886 to->weakreflist = NULL;
887 PyObject_GC_Track(to);
Raymond Hettingerad983e72003-11-12 14:32:26 +0000888done:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000889 Py_XDECREF(it);
890 return (PyObject *)to;
Raymond Hettingerad983e72003-11-12 14:32:26 +0000891}
892
Tal Einatc4bccd32018-09-12 00:49:13 +0300893/*[clinic input]
894@classmethod
895itertools._tee.__new__
896 iterable: object
897 /
898Iterator wrapped to make it copyable.
899[clinic start generated code]*/
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000900
Tal Einatc4bccd32018-09-12 00:49:13 +0300901static PyObject *
902itertools__tee_impl(PyTypeObject *type, PyObject *iterable)
903/*[clinic end generated code: output=b02d3fd26c810c3f input=adc0779d2afe37a2]*/
904{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000905 return tee_fromiterable(iterable);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000906}
907
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000908static int
909tee_clear(teeobject *to)
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000910{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000911 if (to->weakreflist != NULL)
912 PyObject_ClearWeakRefs((PyObject *) to);
913 Py_CLEAR(to->dataobj);
914 return 0;
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000915}
916
917static void
918tee_dealloc(teeobject *to)
919{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000920 PyObject_GC_UnTrack(to);
921 tee_clear(to);
922 PyObject_GC_Del(to);
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000923}
924
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000925static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +0530926tee_reduce(teeobject *to, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000927{
928 return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
929}
930
931static PyObject *
932tee_setstate(teeobject *to, PyObject *state)
933{
934 teedataobject *tdo;
935 int index;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300936 if (!PyTuple_Check(state)) {
937 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000938 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +0300939 }
940 if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {
941 return NULL;
942 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000943 if (index < 0 || index > LINKCELLS) {
944 PyErr_SetString(PyExc_ValueError, "Index out of range");
945 return NULL;
946 }
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +0200947 Py_INCREF(tdo);
Serhiy Storchaka48842712016-04-06 09:45:48 +0300948 Py_XSETREF(to->dataobj, tdo);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000949 to->index = index;
950 Py_RETURN_NONE;
951}
952
Raymond Hettingerad983e72003-11-12 14:32:26 +0000953static PyMethodDef tee_methods[] = {
Raymond Hettingera6a2d442015-08-16 14:38:07 -0700954 {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
955 {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
956 {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000957 {NULL, NULL} /* sentinel */
Raymond Hettingerad983e72003-11-12 14:32:26 +0000958};
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000959
960static PyTypeObject tee_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000961 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000962 "itertools._tee", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000963 sizeof(teeobject), /* tp_basicsize */
964 0, /* tp_itemsize */
965 /* methods */
966 (destructor)tee_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200967 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000968 0, /* tp_getattr */
969 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +0200970 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000971 0, /* tp_repr */
972 0, /* tp_as_number */
973 0, /* tp_as_sequence */
974 0, /* tp_as_mapping */
975 0, /* tp_hash */
976 0, /* tp_call */
977 0, /* tp_str */
978 0, /* tp_getattro */
979 0, /* tp_setattro */
980 0, /* tp_as_buffer */
981 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +0300982 itertools__tee__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000983 (traverseproc)tee_traverse, /* tp_traverse */
984 (inquiry)tee_clear, /* tp_clear */
985 0, /* tp_richcompare */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -0700986 offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +0000987 PyObject_SelfIter, /* tp_iter */
988 (iternextfunc)tee_next, /* tp_iternext */
989 tee_methods, /* tp_methods */
990 0, /* tp_members */
991 0, /* tp_getset */
992 0, /* tp_base */
993 0, /* tp_dict */
994 0, /* tp_descr_get */
995 0, /* tp_descr_set */
996 0, /* tp_dictoffset */
997 0, /* tp_init */
998 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +0300999 itertools__tee, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001000 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6a5b0272003-10-24 08:45:23 +00001001};
1002
Tal Einatc4bccd32018-09-12 00:49:13 +03001003/*[clinic input]
1004itertools.tee
1005 iterable: object
1006 n: Py_ssize_t = 2
1007 /
1008Returns a tuple of n independent iterators.
1009[clinic start generated code]*/
1010
Raymond Hettingerad983e72003-11-12 14:32:26 +00001011static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001012itertools_tee_impl(PyObject *module, PyObject *iterable, Py_ssize_t n)
1013/*[clinic end generated code: output=1c64519cd859c2f0 input=c99a1472c425d66d]*/
Raymond Hettingerad983e72003-11-12 14:32:26 +00001014{
Tal Einatc4bccd32018-09-12 00:49:13 +03001015 Py_ssize_t i;
1016 PyObject *it, *copyable, *copyfunc, *result;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02001017 _Py_IDENTIFIER(__copy__);
Raymond Hettingerad983e72003-11-12 14:32:26 +00001018
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001019 if (n < 0) {
1020 PyErr_SetString(PyExc_ValueError, "n must be >= 0");
1021 return NULL;
1022 }
1023 result = PyTuple_New(n);
1024 if (result == NULL)
1025 return NULL;
1026 if (n == 0)
1027 return result;
1028 it = PyObject_GetIter(iterable);
1029 if (it == NULL) {
1030 Py_DECREF(result);
1031 return NULL;
1032 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001033
Serhiy Storchakaf320be72018-01-25 10:49:40 +02001034 if (_PyObject_LookupAttrId(it, &PyId___copy__, &copyfunc) < 0) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001035 Py_DECREF(it);
1036 Py_DECREF(result);
1037 return NULL;
1038 }
Serhiy Storchakaf320be72018-01-25 10:49:40 +02001039 if (copyfunc != NULL) {
1040 copyable = it;
1041 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001042 else {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001043 copyable = tee_fromiterable(it);
1044 Py_DECREF(it);
1045 if (copyable == NULL) {
1046 Py_DECREF(result);
1047 return NULL;
1048 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001049 copyfunc = _PyObject_GetAttrId(copyable, &PyId___copy__);
1050 if (copyfunc == NULL) {
1051 Py_DECREF(copyable);
1052 Py_DECREF(result);
1053 return NULL;
1054 }
1055 }
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02001056
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001057 PyTuple_SET_ITEM(result, 0, copyable);
1058 for (i = 1; i < n; i++) {
1059 copyable = _PyObject_CallNoArg(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001060 if (copyable == NULL) {
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001061 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001062 Py_DECREF(result);
1063 return NULL;
1064 }
1065 PyTuple_SET_ITEM(result, i, copyable);
1066 }
Serhiy Storchaka1707e402017-11-11 15:51:42 +02001067 Py_DECREF(copyfunc);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001068 return result;
Raymond Hettingerad983e72003-11-12 14:32:26 +00001069}
1070
Raymond Hettingerad983e72003-11-12 14:32:26 +00001071
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001072/* cycle object **************************************************************/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001073
1074typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001075 PyObject_HEAD
1076 PyObject *it;
1077 PyObject *saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001078 Py_ssize_t index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001079 int firstpass;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001080} cycleobject;
1081
Tal Einatc4bccd32018-09-12 00:49:13 +03001082/*[clinic input]
1083@classmethod
1084itertools.cycle.__new__
1085 iterable: object
1086 /
1087Return elements from the iterable until it is exhausted. Then repeat the sequence indefinitely.
1088[clinic start generated code]*/
1089
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001090static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001091itertools_cycle_impl(PyTypeObject *type, PyObject *iterable)
1092/*[clinic end generated code: output=f60e5ec17a45b35c input=9d1d84bcf66e908b]*/
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001093{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001094 PyObject *it;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001095 PyObject *saved;
1096 cycleobject *lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001097
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001098 /* Get iterator. */
1099 it = PyObject_GetIter(iterable);
1100 if (it == NULL)
1101 return NULL;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001102
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001103 saved = PyList_New(0);
1104 if (saved == NULL) {
1105 Py_DECREF(it);
1106 return NULL;
1107 }
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001108
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001109 /* create cycleobject structure */
1110 lz = (cycleobject *)type->tp_alloc(type, 0);
1111 if (lz == NULL) {
1112 Py_DECREF(it);
1113 Py_DECREF(saved);
1114 return NULL;
1115 }
1116 lz->it = it;
1117 lz->saved = saved;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001118 lz->index = 0;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001119 lz->firstpass = 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001120
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001121 return (PyObject *)lz;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001122}
1123
1124static void
1125cycle_dealloc(cycleobject *lz)
1126{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001127 PyObject_GC_UnTrack(lz);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001128 Py_XDECREF(lz->it);
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001129 Py_XDECREF(lz->saved);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001130 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001131}
1132
1133static int
1134cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
1135{
Hai Shi47a23fc2020-06-07 20:05:36 +08001136 Py_VISIT(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001137 Py_VISIT(lz->saved);
1138 return 0;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001139}
1140
1141static PyObject *
1142cycle_next(cycleobject *lz)
1143{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001144 PyObject *item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001145
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001146 if (lz->it != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001147 item = PyIter_Next(lz->it);
1148 if (item != NULL) {
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001149 if (lz->firstpass)
1150 return item;
1151 if (PyList_Append(lz->saved, item)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001152 Py_DECREF(item);
1153 return NULL;
1154 }
1155 return item;
1156 }
Raymond Hettinger98958fe2015-08-15 15:09:30 -07001157 /* Note: StopIteration is already cleared by PyIter_Next() */
1158 if (PyErr_Occurred())
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001159 return NULL;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001160 Py_CLEAR(lz->it);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001161 }
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001162 if (PyList_GET_SIZE(lz->saved) == 0)
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001163 return NULL;
1164 item = PyList_GET_ITEM(lz->saved, lz->index);
1165 lz->index++;
Serhiy Storchakafff9a312017-03-21 08:53:25 +02001166 if (lz->index >= PyList_GET_SIZE(lz->saved))
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001167 lz->index = 0;
1168 Py_INCREF(item);
1169 return item;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001170}
1171
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001172static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301173cycle_reduce(cycleobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001174{
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001175 /* Create a new cycle with the iterator tuple, then set the saved state */
1176 if (lz->it == NULL) {
1177 PyObject *it = PyObject_GetIter(lz->saved);
1178 if (it == NULL)
1179 return NULL;
1180 if (lz->index != 0) {
1181 _Py_IDENTIFIER(__setstate__);
1182 PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
1183 "n", lz->index);
1184 if (res == NULL) {
1185 Py_DECREF(it);
1186 return NULL;
1187 }
1188 Py_DECREF(res);
1189 }
Serhiy Storchaka1f21eaa2019-09-01 12:16:51 +03001190 return Py_BuildValue("O(N)(OO)", Py_TYPE(lz), it, lz->saved, Py_True);
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001191 }
Serhiy Storchaka1f21eaa2019-09-01 12:16:51 +03001192 return Py_BuildValue("O(O)(OO)", Py_TYPE(lz), lz->it, lz->saved,
1193 lz->firstpass ? Py_True : Py_False);
Raymond Hettingerc39786d2015-08-15 15:16:12 -07001194}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001195
1196static PyObject *
1197cycle_setstate(cycleobject *lz, PyObject *state)
1198{
1199 PyObject *saved=NULL;
1200 int firstpass;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001201 if (!PyTuple_Check(state)) {
1202 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001203 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03001204 }
1205 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
1206 return NULL;
1207 }
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001208 Py_INCREF(saved);
Serhiy Storchaka48842712016-04-06 09:45:48 +03001209 Py_XSETREF(lz->saved, saved);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001210 lz->firstpass = firstpass != 0;
Raymond Hettingerca3788c2015-08-16 14:49:24 -07001211 lz->index = 0;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001212 Py_RETURN_NONE;
1213}
1214
1215static PyMethodDef cycle_methods[] = {
1216 {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
1217 reduce_doc},
1218 {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
1219 setstate_doc},
1220 {NULL, NULL} /* sentinel */
1221};
1222
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001223static PyTypeObject cycle_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001224 PyVarObject_HEAD_INIT(NULL, 0)
1225 "itertools.cycle", /* tp_name */
1226 sizeof(cycleobject), /* tp_basicsize */
1227 0, /* tp_itemsize */
1228 /* methods */
1229 (destructor)cycle_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001230 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001231 0, /* tp_getattr */
1232 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001233 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001234 0, /* tp_repr */
1235 0, /* tp_as_number */
1236 0, /* tp_as_sequence */
1237 0, /* tp_as_mapping */
1238 0, /* tp_hash */
1239 0, /* tp_call */
1240 0, /* tp_str */
1241 PyObject_GenericGetAttr, /* tp_getattro */
1242 0, /* tp_setattro */
1243 0, /* tp_as_buffer */
1244 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1245 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001246 itertools_cycle__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001247 (traverseproc)cycle_traverse, /* tp_traverse */
1248 0, /* tp_clear */
1249 0, /* tp_richcompare */
1250 0, /* tp_weaklistoffset */
1251 PyObject_SelfIter, /* tp_iter */
1252 (iternextfunc)cycle_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001253 cycle_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001254 0, /* tp_members */
1255 0, /* tp_getset */
1256 0, /* tp_base */
1257 0, /* tp_dict */
1258 0, /* tp_descr_get */
1259 0, /* tp_descr_set */
1260 0, /* tp_dictoffset */
1261 0, /* tp_init */
1262 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001263 itertools_cycle, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001264 PyObject_GC_Del, /* tp_free */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001265};
1266
1267
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001268/* dropwhile object **********************************************************/
1269
1270typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001271 PyObject_HEAD
1272 PyObject *func;
1273 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001274 long start;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001275} dropwhileobject;
1276
Tal Einatc4bccd32018-09-12 00:49:13 +03001277/*[clinic input]
1278@classmethod
1279itertools.dropwhile.__new__
1280 predicate as func: object
1281 iterable as seq: object
1282 /
1283Drop items from the iterable while predicate(item) is true.
1284
1285Afterwards, return every element until the iterable is exhausted.
1286[clinic start generated code]*/
1287
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001288static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001289itertools_dropwhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1290/*[clinic end generated code: output=92f9d0d89af149e4 input=d39737147c9f0a26]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001291{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001292 PyObject *it;
1293 dropwhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001294
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001295 /* Get iterator. */
1296 it = PyObject_GetIter(seq);
1297 if (it == NULL)
1298 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001299
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001300 /* create dropwhileobject structure */
1301 lz = (dropwhileobject *)type->tp_alloc(type, 0);
1302 if (lz == NULL) {
1303 Py_DECREF(it);
1304 return NULL;
1305 }
1306 Py_INCREF(func);
1307 lz->func = func;
1308 lz->it = it;
1309 lz->start = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001310
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001311 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001312}
1313
1314static void
1315dropwhile_dealloc(dropwhileobject *lz)
1316{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001317 PyObject_GC_UnTrack(lz);
1318 Py_XDECREF(lz->func);
1319 Py_XDECREF(lz->it);
1320 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001321}
1322
1323static int
1324dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1325{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001326 Py_VISIT(lz->it);
1327 Py_VISIT(lz->func);
1328 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001329}
1330
1331static PyObject *
1332dropwhile_next(dropwhileobject *lz)
1333{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001334 PyObject *item, *good;
1335 PyObject *it = lz->it;
1336 long ok;
1337 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001338
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001339 iternext = *Py_TYPE(it)->tp_iternext;
1340 for (;;) {
1341 item = iternext(it);
1342 if (item == NULL)
1343 return NULL;
1344 if (lz->start == 1)
1345 return item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001346
Petr Viktorinffd97532020-02-11 17:46:57 +01001347 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001348 if (good == NULL) {
1349 Py_DECREF(item);
1350 return NULL;
1351 }
1352 ok = PyObject_IsTrue(good);
1353 Py_DECREF(good);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001354 if (ok == 0) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001355 lz->start = 1;
1356 return item;
1357 }
1358 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001359 if (ok < 0)
1360 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001361 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001362}
1363
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001364static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301365dropwhile_reduce(dropwhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001366{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001367 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 +00001368}
1369
1370static PyObject *
1371dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1372{
1373 int start = PyObject_IsTrue(state);
Antoine Pitrou721738f2012-08-15 23:20:39 +02001374 if (start < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001375 return NULL;
1376 lz->start = start;
1377 Py_RETURN_NONE;
1378}
1379
1380static PyMethodDef dropwhile_methods[] = {
1381 {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1382 reduce_doc},
1383 {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1384 setstate_doc},
1385 {NULL, NULL} /* sentinel */
1386};
1387
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001388static PyTypeObject dropwhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001389 PyVarObject_HEAD_INIT(NULL, 0)
1390 "itertools.dropwhile", /* tp_name */
1391 sizeof(dropwhileobject), /* tp_basicsize */
1392 0, /* tp_itemsize */
1393 /* methods */
1394 (destructor)dropwhile_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001395 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001396 0, /* tp_getattr */
1397 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001398 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001399 0, /* tp_repr */
1400 0, /* tp_as_number */
1401 0, /* tp_as_sequence */
1402 0, /* tp_as_mapping */
1403 0, /* tp_hash */
1404 0, /* tp_call */
1405 0, /* tp_str */
1406 PyObject_GenericGetAttr, /* tp_getattro */
1407 0, /* tp_setattro */
1408 0, /* tp_as_buffer */
1409 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1410 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001411 itertools_dropwhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001412 (traverseproc)dropwhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001413 0, /* tp_clear */
1414 0, /* tp_richcompare */
1415 0, /* tp_weaklistoffset */
1416 PyObject_SelfIter, /* tp_iter */
1417 (iternextfunc)dropwhile_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001418 dropwhile_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001419 0, /* tp_members */
1420 0, /* tp_getset */
1421 0, /* tp_base */
1422 0, /* tp_dict */
1423 0, /* tp_descr_get */
1424 0, /* tp_descr_set */
1425 0, /* tp_dictoffset */
1426 0, /* tp_init */
1427 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001428 itertools_dropwhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001429 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001430};
1431
1432
1433/* takewhile object **********************************************************/
1434
1435typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001436 PyObject_HEAD
1437 PyObject *func;
1438 PyObject *it;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001439 long stop;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001440} takewhileobject;
1441
Tal Einatc4bccd32018-09-12 00:49:13 +03001442/*[clinic input]
1443@classmethod
1444itertools.takewhile.__new__
1445 predicate as func: object
1446 iterable as seq: object
1447 /
1448Return successive entries from an iterable as long as the predicate evaluates to true for each entry.
1449[clinic start generated code]*/
1450
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001451static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001452itertools_takewhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1453/*[clinic end generated code: output=bb179ea7864e2ef6 input=ba5255f7519aa119]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001454{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001455 PyObject *it;
1456 takewhileobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001457
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001458 /* Get iterator. */
1459 it = PyObject_GetIter(seq);
1460 if (it == NULL)
1461 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001462
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001463 /* create takewhileobject structure */
1464 lz = (takewhileobject *)type->tp_alloc(type, 0);
1465 if (lz == NULL) {
1466 Py_DECREF(it);
1467 return NULL;
1468 }
1469 Py_INCREF(func);
1470 lz->func = func;
1471 lz->it = it;
1472 lz->stop = 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001473
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001474 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001475}
1476
1477static void
1478takewhile_dealloc(takewhileobject *lz)
1479{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001480 PyObject_GC_UnTrack(lz);
1481 Py_XDECREF(lz->func);
1482 Py_XDECREF(lz->it);
1483 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001484}
1485
1486static int
1487takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1488{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001489 Py_VISIT(lz->it);
1490 Py_VISIT(lz->func);
1491 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001492}
1493
1494static PyObject *
1495takewhile_next(takewhileobject *lz)
1496{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001497 PyObject *item, *good;
1498 PyObject *it = lz->it;
1499 long ok;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001500
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001501 if (lz->stop == 1)
1502 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001503
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001504 item = (*Py_TYPE(it)->tp_iternext)(it);
1505 if (item == NULL)
1506 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001507
Petr Viktorinffd97532020-02-11 17:46:57 +01001508 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001509 if (good == NULL) {
1510 Py_DECREF(item);
1511 return NULL;
1512 }
1513 ok = PyObject_IsTrue(good);
1514 Py_DECREF(good);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001515 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001516 return item;
1517 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02001518 if (ok == 0)
1519 lz->stop = 1;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001520 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001521}
1522
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001523static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301524takewhile_reduce(takewhileobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001525{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001526 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 +00001527}
1528
1529static PyObject *
1530takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1531{
1532 int stop = PyObject_IsTrue(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001533
Antoine Pitrou721738f2012-08-15 23:20:39 +02001534 if (stop < 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001535 return NULL;
1536 lz->stop = stop;
1537 Py_RETURN_NONE;
1538}
1539
1540static PyMethodDef takewhile_reduce_methods[] = {
1541 {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1542 reduce_doc},
1543 {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1544 setstate_doc},
1545 {NULL, NULL} /* sentinel */
1546};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001547
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001548static PyTypeObject takewhile_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001549 PyVarObject_HEAD_INIT(NULL, 0)
1550 "itertools.takewhile", /* tp_name */
1551 sizeof(takewhileobject), /* tp_basicsize */
1552 0, /* tp_itemsize */
1553 /* methods */
1554 (destructor)takewhile_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001555 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001556 0, /* tp_getattr */
1557 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001558 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001559 0, /* tp_repr */
1560 0, /* tp_as_number */
1561 0, /* tp_as_sequence */
1562 0, /* tp_as_mapping */
1563 0, /* tp_hash */
1564 0, /* tp_call */
1565 0, /* tp_str */
1566 PyObject_GenericGetAttr, /* tp_getattro */
1567 0, /* tp_setattro */
1568 0, /* tp_as_buffer */
1569 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1570 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001571 itertools_takewhile__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07001572 (traverseproc)takewhile_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001573 0, /* tp_clear */
1574 0, /* tp_richcompare */
1575 0, /* tp_weaklistoffset */
1576 PyObject_SelfIter, /* tp_iter */
1577 (iternextfunc)takewhile_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001578 takewhile_reduce_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001579 0, /* tp_members */
1580 0, /* tp_getset */
1581 0, /* tp_base */
1582 0, /* tp_dict */
1583 0, /* tp_descr_get */
1584 0, /* tp_descr_set */
1585 0, /* tp_dictoffset */
1586 0, /* tp_init */
1587 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001588 itertools_takewhile, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001589 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001590};
1591
1592
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001593/* islice object *************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001594
1595typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001596 PyObject_HEAD
1597 PyObject *it;
1598 Py_ssize_t next;
1599 Py_ssize_t stop;
1600 Py_ssize_t step;
1601 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001602} isliceobject;
1603
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001604static PyTypeObject islice_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001605
1606static PyObject *
1607islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1608{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001609 PyObject *seq;
1610 Py_ssize_t start=0, stop=-1, step=1;
1611 PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1612 Py_ssize_t numargs;
1613 isliceobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001614
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03001615 if (type == &islice_type && !_PyArg_NoKeywords("islice", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001616 return NULL;
Georg Brandl02c42872005-08-26 06:42:30 +00001617
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001618 if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1619 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001620
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001621 numargs = PyTuple_Size(args);
1622 if (numargs == 2) {
1623 if (a1 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001624 stop = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001625 if (stop == -1) {
1626 if (PyErr_Occurred())
1627 PyErr_Clear();
1628 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001629 "Stop argument for islice() must be None or "
1630 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001631 return NULL;
1632 }
1633 }
1634 } else {
1635 if (a1 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001636 start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001637 if (start == -1 && PyErr_Occurred())
1638 PyErr_Clear();
1639 if (a2 != Py_None) {
Will Roberts0ecdc522017-06-08 08:03:04 +02001640 stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001641 if (stop == -1) {
1642 if (PyErr_Occurred())
1643 PyErr_Clear();
1644 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001645 "Stop argument for islice() must be None or "
1646 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001647 return NULL;
1648 }
1649 }
1650 }
1651 if (start<0 || stop<-1) {
1652 PyErr_SetString(PyExc_ValueError,
Raymond Hettingera6a2d442015-08-16 14:38:07 -07001653 "Indices for islice() must be None or "
1654 "an integer: 0 <= x <= sys.maxsize.");
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001655 return NULL;
1656 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001657
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001658 if (a3 != NULL) {
1659 if (a3 != Py_None)
Will Roberts0ecdc522017-06-08 08:03:04 +02001660 step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001661 if (step == -1 && PyErr_Occurred())
1662 PyErr_Clear();
1663 }
1664 if (step<1) {
1665 PyErr_SetString(PyExc_ValueError,
1666 "Step for islice() must be a positive integer or None.");
1667 return NULL;
1668 }
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001669
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001670 /* Get iterator. */
1671 it = PyObject_GetIter(seq);
1672 if (it == NULL)
1673 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001674
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001675 /* create isliceobject structure */
1676 lz = (isliceobject *)type->tp_alloc(type, 0);
1677 if (lz == NULL) {
1678 Py_DECREF(it);
1679 return NULL;
1680 }
1681 lz->it = it;
1682 lz->next = start;
1683 lz->stop = stop;
1684 lz->step = step;
1685 lz->cnt = 0L;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001686
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001687 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001688}
1689
1690static void
1691islice_dealloc(isliceobject *lz)
1692{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001693 PyObject_GC_UnTrack(lz);
1694 Py_XDECREF(lz->it);
1695 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001696}
1697
1698static int
1699islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1700{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001701 Py_VISIT(lz->it);
1702 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001703}
1704
1705static PyObject *
1706islice_next(isliceobject *lz)
1707{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001708 PyObject *item;
1709 PyObject *it = lz->it;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001710 Py_ssize_t stop = lz->stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001711 Py_ssize_t oldnext;
1712 PyObject *(*iternext)(PyObject *);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001713
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001714 if (it == NULL)
1715 return NULL;
1716
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001717 iternext = *Py_TYPE(it)->tp_iternext;
1718 while (lz->cnt < lz->next) {
1719 item = iternext(it);
1720 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001721 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001722 Py_DECREF(item);
1723 lz->cnt++;
1724 }
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001725 if (stop != -1 && lz->cnt >= stop)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001726 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001727 item = iternext(it);
1728 if (item == NULL)
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001729 goto empty;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001730 lz->cnt++;
1731 oldnext = lz->next;
Mark Dickinsonb2f6bc72011-09-24 08:56:09 +01001732 /* The (size_t) cast below avoids the danger of undefined
1733 behaviour from signed integer overflow. */
1734 lz->next += (size_t)lz->step;
Raymond Hettinger69b34bf2010-11-30 02:49:29 +00001735 if (lz->next < oldnext || (stop != -1 && lz->next > stop))
1736 lz->next = stop;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001737 return item;
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001738
1739empty:
1740 Py_CLEAR(lz->it);
1741 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001742}
1743
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001744static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301745islice_reduce(isliceobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001746{
1747 /* When unpickled, generate a new object with the same bounds,
1748 * then 'setstate' with the next and count
1749 */
1750 PyObject *stop;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001751
Antoine Pitrou26f82ef2014-04-29 12:13:46 +02001752 if (lz->it == NULL) {
1753 PyObject *empty_list;
1754 PyObject *empty_it;
1755 empty_list = PyList_New(0);
1756 if (empty_list == NULL)
1757 return NULL;
1758 empty_it = PyObject_GetIter(empty_list);
1759 Py_DECREF(empty_list);
1760 if (empty_it == NULL)
1761 return NULL;
1762 return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1763 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001764 if (lz->stop == -1) {
1765 stop = Py_None;
1766 Py_INCREF(stop);
1767 } else {
1768 stop = PyLong_FromSsize_t(lz->stop);
1769 if (stop == NULL)
1770 return NULL;
1771 }
1772 return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1773 lz->it, lz->next, stop, lz->step,
1774 lz->cnt);
1775}
1776
1777static PyObject *
1778islice_setstate(isliceobject *lz, PyObject *state)
1779{
1780 Py_ssize_t cnt = PyLong_AsSsize_t(state);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001781
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001782 if (cnt == -1 && PyErr_Occurred())
1783 return NULL;
1784 lz->cnt = cnt;
1785 Py_RETURN_NONE;
1786}
1787
1788static PyMethodDef islice_methods[] = {
1789 {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1790 reduce_doc},
1791 {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1792 setstate_doc},
1793 {NULL, NULL} /* sentinel */
1794};
1795
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001796PyDoc_STRVAR(islice_doc,
Andrew Kuchlingda30acf2013-06-22 19:20:54 -04001797"islice(iterable, stop) --> islice object\n\
1798islice(iterable, start, stop[, step]) --> islice object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001799\n\
1800Return an iterator whose next() method returns selected values from an\n\
1801iterable. If start is specified, will skip all preceding elements;\n\
1802otherwise, start defaults to zero. Step defaults to one. If\n\
oldkaa0735f2018-02-02 16:52:55 +08001803specified as another value, step determines how many values are\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001804skipped between successive calls. Works like a slice() on a list\n\
1805but returns an iterator.");
1806
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001807static PyTypeObject islice_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001808 PyVarObject_HEAD_INIT(NULL, 0)
1809 "itertools.islice", /* tp_name */
1810 sizeof(isliceobject), /* tp_basicsize */
1811 0, /* tp_itemsize */
1812 /* methods */
1813 (destructor)islice_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001814 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001815 0, /* tp_getattr */
1816 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001817 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001818 0, /* tp_repr */
1819 0, /* tp_as_number */
1820 0, /* tp_as_sequence */
1821 0, /* tp_as_mapping */
1822 0, /* tp_hash */
1823 0, /* tp_call */
1824 0, /* tp_str */
1825 PyObject_GenericGetAttr, /* tp_getattro */
1826 0, /* tp_setattro */
1827 0, /* tp_as_buffer */
1828 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1829 Py_TPFLAGS_BASETYPE, /* tp_flags */
1830 islice_doc, /* tp_doc */
1831 (traverseproc)islice_traverse, /* tp_traverse */
1832 0, /* tp_clear */
1833 0, /* tp_richcompare */
1834 0, /* tp_weaklistoffset */
1835 PyObject_SelfIter, /* tp_iter */
1836 (iternextfunc)islice_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001837 islice_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001838 0, /* tp_members */
1839 0, /* tp_getset */
1840 0, /* tp_base */
1841 0, /* tp_dict */
1842 0, /* tp_descr_get */
1843 0, /* tp_descr_set */
1844 0, /* tp_dictoffset */
1845 0, /* tp_init */
1846 0, /* tp_alloc */
1847 islice_new, /* tp_new */
1848 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001849};
1850
1851
1852/* starmap object ************************************************************/
1853
1854typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001855 PyObject_HEAD
1856 PyObject *func;
1857 PyObject *it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001858} starmapobject;
1859
Tal Einatc4bccd32018-09-12 00:49:13 +03001860/*[clinic input]
1861@classmethod
1862itertools.starmap.__new__
1863 function as func: object
1864 iterable as seq: object
1865 /
1866Return an iterator whose values are returned from the function evaluated with an argument tuple taken from the given sequence.
1867[clinic start generated code]*/
1868
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001869static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03001870itertools_starmap_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1871/*[clinic end generated code: output=79eeb81d452c6e8d input=844766df6a0d4dad]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001872{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001873 PyObject *it;
1874 starmapobject *lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001875
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001876 /* Get iterator. */
1877 it = PyObject_GetIter(seq);
1878 if (it == NULL)
1879 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001880
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001881 /* create starmapobject structure */
1882 lz = (starmapobject *)type->tp_alloc(type, 0);
1883 if (lz == NULL) {
1884 Py_DECREF(it);
1885 return NULL;
1886 }
1887 Py_INCREF(func);
1888 lz->func = func;
1889 lz->it = it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001890
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001891 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001892}
1893
1894static void
1895starmap_dealloc(starmapobject *lz)
1896{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001897 PyObject_GC_UnTrack(lz);
1898 Py_XDECREF(lz->func);
1899 Py_XDECREF(lz->it);
1900 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001901}
1902
1903static int
1904starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1905{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001906 Py_VISIT(lz->it);
1907 Py_VISIT(lz->func);
1908 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001909}
1910
1911static PyObject *
1912starmap_next(starmapobject *lz)
1913{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001914 PyObject *args;
1915 PyObject *result;
1916 PyObject *it = lz->it;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001917
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001918 args = (*Py_TYPE(it)->tp_iternext)(it);
1919 if (args == NULL)
1920 return NULL;
1921 if (!PyTuple_CheckExact(args)) {
1922 PyObject *newargs = PySequence_Tuple(args);
1923 Py_DECREF(args);
1924 if (newargs == NULL)
1925 return NULL;
1926 args = newargs;
1927 }
1928 result = PyObject_Call(lz->func, args, NULL);
1929 Py_DECREF(args);
1930 return result;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001931}
1932
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001933static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05301934starmap_reduce(starmapobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001935{
1936 /* Just pickle the iterator */
1937 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1938}
1939
1940static PyMethodDef starmap_methods[] = {
1941 {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1942 reduce_doc},
1943 {NULL, NULL} /* sentinel */
1944};
1945
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001946static PyTypeObject starmap_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001947 PyVarObject_HEAD_INIT(NULL, 0)
1948 "itertools.starmap", /* tp_name */
1949 sizeof(starmapobject), /* tp_basicsize */
1950 0, /* tp_itemsize */
1951 /* methods */
1952 (destructor)starmap_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001953 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001954 0, /* tp_getattr */
1955 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02001956 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001957 0, /* tp_repr */
1958 0, /* tp_as_number */
1959 0, /* tp_as_sequence */
1960 0, /* tp_as_mapping */
1961 0, /* tp_hash */
1962 0, /* tp_call */
1963 0, /* tp_str */
1964 PyObject_GenericGetAttr, /* tp_getattro */
1965 0, /* tp_setattro */
1966 0, /* tp_as_buffer */
1967 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1968 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03001969 itertools_starmap__doc__, /* tp_doc */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001970 (traverseproc)starmap_traverse, /* tp_traverse */
1971 0, /* tp_clear */
1972 0, /* tp_richcompare */
1973 0, /* tp_weaklistoffset */
1974 PyObject_SelfIter, /* tp_iter */
1975 (iternextfunc)starmap_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00001976 starmap_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001977 0, /* tp_members */
1978 0, /* tp_getset */
1979 0, /* tp_base */
1980 0, /* tp_dict */
1981 0, /* tp_descr_get */
1982 0, /* tp_descr_set */
1983 0, /* tp_dictoffset */
1984 0, /* tp_init */
1985 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03001986 itertools_starmap, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001987 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001988};
1989
1990
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07001991/* chain object **************************************************************/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001992
1993typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00001994 PyObject_HEAD
1995 PyObject *source; /* Iterator over input iterables */
1996 PyObject *active; /* Currently running input iterator */
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00001997} chainobject;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001998
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00001999static PyTypeObject chain_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002000
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002001static PyObject *
Christian Heimesf16baeb2008-02-29 14:57:44 +00002002chain_new_internal(PyTypeObject *type, PyObject *source)
2003{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002004 chainobject *lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00002005
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002006 lz = (chainobject *)type->tp_alloc(type, 0);
2007 if (lz == NULL) {
2008 Py_DECREF(source);
2009 return NULL;
2010 }
2011
2012 lz->source = source;
2013 lz->active = NULL;
2014 return (PyObject *)lz;
Christian Heimesf16baeb2008-02-29 14:57:44 +00002015}
2016
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002017static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002018chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002019{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002020 PyObject *source;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002021
Serhiy Storchaka6cca5c82017-06-08 14:41:19 +03002022 if (type == &chain_type && !_PyArg_NoKeywords("chain", kwds))
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002023 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002024
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002025 source = PyObject_GetIter(args);
2026 if (source == NULL)
2027 return NULL;
2028
2029 return chain_new_internal(type, source);
Christian Heimesf16baeb2008-02-29 14:57:44 +00002030}
2031
Tal Einatc4bccd32018-09-12 00:49:13 +03002032/*[clinic input]
2033@classmethod
2034itertools.chain.from_iterable
2035 iterable as arg: object
2036 /
2037Alternative chain() constructor taking a single iterable argument that evaluates lazily.
2038[clinic start generated code]*/
2039
Christian Heimesf16baeb2008-02-29 14:57:44 +00002040static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002041itertools_chain_from_iterable(PyTypeObject *type, PyObject *arg)
2042/*[clinic end generated code: output=667ae7a7f7b68654 input=72c39e3a2ca3be85]*/
Christian Heimesf16baeb2008-02-29 14:57:44 +00002043{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002044 PyObject *source;
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002045
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002046 source = PyObject_GetIter(arg);
2047 if (source == NULL)
2048 return NULL;
2049
2050 return chain_new_internal(type, source);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002051}
2052
2053static void
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002054chain_dealloc(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002055{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002056 PyObject_GC_UnTrack(lz);
2057 Py_XDECREF(lz->active);
2058 Py_XDECREF(lz->source);
2059 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002060}
2061
Raymond Hettinger2012f172003-02-07 05:32:58 +00002062static int
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002063chain_traverse(chainobject *lz, visitproc visit, void *arg)
Raymond Hettinger2012f172003-02-07 05:32:58 +00002064{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002065 Py_VISIT(lz->source);
2066 Py_VISIT(lz->active);
2067 return 0;
Raymond Hettinger2012f172003-02-07 05:32:58 +00002068}
2069
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002070static PyObject *
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002071chain_next(chainobject *lz)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002072{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002073 PyObject *item;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002074
T. Wouters5466d4a2017-03-30 09:58:35 -07002075 /* lz->source is the iterator of iterables. If it's NULL, we've already
2076 * consumed them all. lz->active is the current iterator. If it's NULL,
2077 * we should grab a new one from lz->source. */
2078 while (lz->source != NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002079 if (lz->active == NULL) {
T. Wouters5466d4a2017-03-30 09:58:35 -07002080 PyObject *iterable = PyIter_Next(lz->source);
2081 if (iterable == NULL) {
2082 Py_CLEAR(lz->source);
2083 return NULL; /* no more input sources */
2084 }
2085 lz->active = PyObject_GetIter(iterable);
2086 Py_DECREF(iterable);
2087 if (lz->active == NULL) {
2088 Py_CLEAR(lz->source);
2089 return NULL; /* input not iterable */
2090 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002091 }
T. Wouters5466d4a2017-03-30 09:58:35 -07002092 item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
2093 if (item != NULL)
2094 return item;
2095 if (PyErr_Occurred()) {
2096 if (PyErr_ExceptionMatches(PyExc_StopIteration))
2097 PyErr_Clear();
2098 else
2099 return NULL; /* input raised an exception */
2100 }
2101 /* lz->active is consumed, try with the next iterable. */
2102 Py_CLEAR(lz->active);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002103 }
T. Wouters5466d4a2017-03-30 09:58:35 -07002104 /* Everything had been consumed already. */
2105 return NULL;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002106}
2107
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002108static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302109chain_reduce(chainobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002110{
2111 if (lz->source) {
2112 /* we can't pickle function objects (itertools.from_iterable) so
2113 * we must use setstate to replace the iterable. One day we
2114 * will fix pickling of functions
2115 */
2116 if (lz->active) {
2117 return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
2118 } else {
2119 return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
2120 }
2121 } else {
2122 return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
2123 }
2124 return NULL;
2125}
2126
2127static PyObject *
2128chain_setstate(chainobject *lz, PyObject *state)
2129{
2130 PyObject *source, *active=NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002131
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002132 if (!PyTuple_Check(state)) {
2133 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002134 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03002135 }
2136 if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
2137 return NULL;
2138 }
2139 if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
2140 PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
2141 return NULL;
2142 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002143
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002144 Py_INCREF(source);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002145 Py_XSETREF(lz->source, source);
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02002146 Py_XINCREF(active);
Serhiy Storchaka48842712016-04-06 09:45:48 +03002147 Py_XSETREF(lz->active, active);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002148 Py_RETURN_NONE;
2149}
2150
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002151PyDoc_STRVAR(chain_doc,
2152"chain(*iterables) --> chain object\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002153\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00002154Return a chain object whose .__next__() method returns elements from the\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00002155first iterable until it is exhausted, then elements from the next\n\
2156iterable, until all of the iterables are exhausted.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002157
Christian Heimesf16baeb2008-02-29 14:57:44 +00002158static PyMethodDef chain_methods[] = {
Tal Einatc4bccd32018-09-12 00:49:13 +03002159 ITERTOOLS_CHAIN_FROM_ITERABLE_METHODDEF
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002160 {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
2161 reduce_doc},
2162 {"__setstate__", (PyCFunction)chain_setstate, METH_O,
2163 setstate_doc},
Ethan Smitha8403d02020-04-09 20:28:08 -07002164 {"__class_getitem__", (PyCFunction)Py_GenericAlias,
2165 METH_O|METH_CLASS, PyDoc_STR("See PEP 585")},
Raymond Hettingera6a2d442015-08-16 14:38:07 -07002166 {NULL, NULL} /* sentinel */
Christian Heimesf16baeb2008-02-29 14:57:44 +00002167};
2168
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00002169static PyTypeObject chain_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002170 PyVarObject_HEAD_INIT(NULL, 0)
2171 "itertools.chain", /* tp_name */
2172 sizeof(chainobject), /* tp_basicsize */
2173 0, /* tp_itemsize */
2174 /* methods */
2175 (destructor)chain_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002176 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002177 0, /* tp_getattr */
2178 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002179 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002180 0, /* tp_repr */
2181 0, /* tp_as_number */
2182 0, /* tp_as_sequence */
2183 0, /* tp_as_mapping */
2184 0, /* tp_hash */
2185 0, /* tp_call */
2186 0, /* tp_str */
2187 PyObject_GenericGetAttr, /* tp_getattro */
2188 0, /* tp_setattro */
2189 0, /* tp_as_buffer */
2190 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2191 Py_TPFLAGS_BASETYPE, /* tp_flags */
2192 chain_doc, /* tp_doc */
2193 (traverseproc)chain_traverse, /* tp_traverse */
2194 0, /* tp_clear */
2195 0, /* tp_richcompare */
2196 0, /* tp_weaklistoffset */
2197 PyObject_SelfIter, /* tp_iter */
2198 (iternextfunc)chain_next, /* tp_iternext */
2199 chain_methods, /* tp_methods */
2200 0, /* tp_members */
2201 0, /* tp_getset */
2202 0, /* tp_base */
2203 0, /* tp_dict */
2204 0, /* tp_descr_get */
2205 0, /* tp_descr_set */
2206 0, /* tp_dictoffset */
2207 0, /* tp_init */
2208 0, /* tp_alloc */
2209 chain_new, /* tp_new */
2210 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00002211};
2212
2213
Christian Heimesc3f30c42008-02-22 16:37:40 +00002214/* product object ************************************************************/
2215
2216typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002217 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002218 PyObject *pools; /* tuple of pool tuples */
2219 Py_ssize_t *indices; /* one index per pool */
2220 PyObject *result; /* most recently returned result tuple */
2221 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002222} productobject;
2223
2224static PyTypeObject product_type;
2225
2226static PyObject *
2227product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2228{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002229 productobject *lz;
2230 Py_ssize_t nargs, npools, repeat=1;
2231 PyObject *pools = NULL;
2232 Py_ssize_t *indices = NULL;
2233 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002234
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002235 if (kwds != NULL) {
2236 char *kwlist[] = {"repeat", 0};
2237 PyObject *tmpargs = PyTuple_New(0);
2238 if (tmpargs == NULL)
2239 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002240 if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2241 kwlist, &repeat)) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002242 Py_DECREF(tmpargs);
2243 return NULL;
2244 }
2245 Py_DECREF(tmpargs);
2246 if (repeat < 0) {
2247 PyErr_SetString(PyExc_ValueError,
2248 "repeat argument cannot be negative");
2249 return NULL;
2250 }
2251 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002252
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002253 assert(PyTuple_CheckExact(args));
2254 if (repeat == 0) {
2255 nargs = 0;
2256 } else {
2257 nargs = PyTuple_GET_SIZE(args);
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002258 if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
Benjamin Peterson0eaabf12015-02-01 21:34:07 -05002259 PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2260 return NULL;
2261 }
2262 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002263 npools = nargs * repeat;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002264
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002265 indices = PyMem_New(Py_ssize_t, npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002266 if (indices == NULL) {
2267 PyErr_NoMemory();
2268 goto error;
2269 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002270
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002271 pools = PyTuple_New(npools);
2272 if (pools == NULL)
2273 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002274
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002275 for (i=0; i < nargs ; ++i) {
2276 PyObject *item = PyTuple_GET_ITEM(args, i);
2277 PyObject *pool = PySequence_Tuple(item);
2278 if (pool == NULL)
2279 goto error;
2280 PyTuple_SET_ITEM(pools, i, pool);
2281 indices[i] = 0;
2282 }
2283 for ( ; i < npools; ++i) {
2284 PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2285 Py_INCREF(pool);
2286 PyTuple_SET_ITEM(pools, i, pool);
2287 indices[i] = 0;
2288 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002289
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002290 /* create productobject structure */
2291 lz = (productobject *)type->tp_alloc(type, 0);
2292 if (lz == NULL)
2293 goto error;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002294
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002295 lz->pools = pools;
2296 lz->indices = indices;
2297 lz->result = NULL;
2298 lz->stopped = 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002299
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002300 return (PyObject *)lz;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002301
2302error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002303 if (indices != NULL)
2304 PyMem_Free(indices);
2305 Py_XDECREF(pools);
2306 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002307}
2308
2309static void
2310product_dealloc(productobject *lz)
2311{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002312 PyObject_GC_UnTrack(lz);
2313 Py_XDECREF(lz->pools);
2314 Py_XDECREF(lz->result);
2315 if (lz->indices != NULL)
2316 PyMem_Free(lz->indices);
2317 Py_TYPE(lz)->tp_free(lz);
Christian Heimesc3f30c42008-02-22 16:37:40 +00002318}
2319
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002320static PyObject *
2321product_sizeof(productobject *lz, void *unused)
2322{
2323 Py_ssize_t res;
2324
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002325 res = _PyObject_SIZE(Py_TYPE(lz));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002326 res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2327 return PyLong_FromSsize_t(res);
2328}
2329
2330PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2331
Christian Heimesc3f30c42008-02-22 16:37:40 +00002332static int
2333product_traverse(productobject *lz, visitproc visit, void *arg)
2334{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002335 Py_VISIT(lz->pools);
2336 Py_VISIT(lz->result);
2337 return 0;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002338}
2339
2340static PyObject *
2341product_next(productobject *lz)
2342{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002343 PyObject *pool;
2344 PyObject *elem;
2345 PyObject *oldelem;
2346 PyObject *pools = lz->pools;
2347 PyObject *result = lz->result;
2348 Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2349 Py_ssize_t i;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002350
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002351 if (lz->stopped)
2352 return NULL;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002353
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002354 if (result == NULL) {
2355 /* On the first pass, return an initial tuple filled with the
2356 first element from each pool. */
2357 result = PyTuple_New(npools);
2358 if (result == NULL)
2359 goto empty;
2360 lz->result = result;
2361 for (i=0; i < npools; i++) {
2362 pool = PyTuple_GET_ITEM(pools, i);
2363 if (PyTuple_GET_SIZE(pool) == 0)
2364 goto empty;
2365 elem = PyTuple_GET_ITEM(pool, 0);
2366 Py_INCREF(elem);
2367 PyTuple_SET_ITEM(result, i, elem);
2368 }
2369 } else {
2370 Py_ssize_t *indices = lz->indices;
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002371
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002372 /* Copy the previous result tuple or re-use it if available */
2373 if (Py_REFCNT(result) > 1) {
2374 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002375 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), npools);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002376 if (result == NULL)
2377 goto empty;
2378 lz->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002379 Py_DECREF(old_result);
2380 }
2381 /* Now, we've got the only copy so we can update it in-place */
2382 assert (npools==0 || Py_REFCNT(result) == 1);
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002383
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002384 /* Update the pool indices right-to-left. Only advance to the
2385 next pool when the previous one rolls-over */
2386 for (i=npools-1 ; i >= 0 ; i--) {
2387 pool = PyTuple_GET_ITEM(pools, i);
2388 indices[i]++;
2389 if (indices[i] == PyTuple_GET_SIZE(pool)) {
2390 /* Roll-over and advance to next pool */
2391 indices[i] = 0;
2392 elem = PyTuple_GET_ITEM(pool, 0);
2393 Py_INCREF(elem);
2394 oldelem = PyTuple_GET_ITEM(result, i);
2395 PyTuple_SET_ITEM(result, i, elem);
2396 Py_DECREF(oldelem);
2397 } else {
2398 /* No rollover. Just increment and stop here. */
2399 elem = PyTuple_GET_ITEM(pool, indices[i]);
2400 Py_INCREF(elem);
2401 oldelem = PyTuple_GET_ITEM(result, i);
2402 PyTuple_SET_ITEM(result, i, elem);
2403 Py_DECREF(oldelem);
2404 break;
2405 }
2406 }
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002407
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002408 /* If i is negative, then the indices have all rolled-over
2409 and we're done. */
2410 if (i < 0)
2411 goto empty;
2412 }
Christian Heimesc3f30c42008-02-22 16:37:40 +00002413
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002414 Py_INCREF(result);
2415 return result;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002416
2417empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002418 lz->stopped = 1;
2419 return NULL;
Christian Heimesc3f30c42008-02-22 16:37:40 +00002420}
2421
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002422static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302423product_reduce(productobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002424{
2425 if (lz->stopped) {
2426 return Py_BuildValue("O(())", Py_TYPE(lz));
2427 } else if (lz->result == NULL) {
2428 return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2429 } else {
2430 PyObject *indices;
2431 Py_ssize_t n, i;
2432
2433 /* we must pickle the indices use them for setstate, and
2434 * additionally indicate that the iterator has started
2435 */
2436 n = PyTuple_GET_SIZE(lz->pools);
2437 indices = PyTuple_New(n);
2438 if (indices == NULL)
2439 return NULL;
2440 for (i=0; i<n; i++){
2441 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2442 if (!index) {
2443 Py_DECREF(indices);
2444 return NULL;
2445 }
2446 PyTuple_SET_ITEM(indices, i, index);
2447 }
2448 return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2449 }
2450}
2451
2452static PyObject *
2453product_setstate(productobject *lz, PyObject *state)
2454{
2455 PyObject *result;
2456 Py_ssize_t n, i;
2457
2458 n = PyTuple_GET_SIZE(lz->pools);
2459 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2460 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2461 return NULL;
2462 }
2463 for (i=0; i<n; i++)
2464 {
2465 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2466 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002467 PyObject* pool;
2468 Py_ssize_t poolsize;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002469 if (index < 0 && PyErr_Occurred())
2470 return NULL; /* not an integer */
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002471 pool = PyTuple_GET_ITEM(lz->pools, i);
2472 poolsize = PyTuple_GET_SIZE(pool);
2473 if (poolsize == 0) {
2474 lz->stopped = 1;
2475 Py_RETURN_NONE;
2476 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002477 /* clamp the index */
2478 if (index < 0)
2479 index = 0;
Kristján Valur Jónsson102764a2015-09-12 15:20:54 +00002480 else if (index > poolsize-1)
2481 index = poolsize-1;
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002482 lz->indices[i] = index;
2483 }
2484
2485 result = PyTuple_New(n);
2486 if (!result)
2487 return NULL;
2488 for (i=0; i<n; i++) {
2489 PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2490 PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2491 Py_INCREF(element);
2492 PyTuple_SET_ITEM(result, i, element);
2493 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03002494 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002495 Py_RETURN_NONE;
2496}
2497
2498static PyMethodDef product_methods[] = {
2499 {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2500 reduce_doc},
2501 {"__setstate__", (PyCFunction)product_setstate, METH_O,
2502 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002503 {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2504 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002505 {NULL, NULL} /* sentinel */
2506};
2507
Christian Heimesc3f30c42008-02-22 16:37:40 +00002508PyDoc_STRVAR(product_doc,
Andrew Kuchling446a39f2013-06-22 19:04:11 -04002509"product(*iterables, repeat=1) --> product object\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002510\n\
Christian Heimes90c3d9b2008-02-23 13:18:03 +00002511Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002512For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2513The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2514cycle in a manner similar to an odometer (with the rightmost element changing\n\
2515on every iteration).\n\n\
Benjamin Petersona86f2c02009-02-10 02:41:10 +00002516To compute the product of an iterable with itself, specify the number\n\
2517of repetitions with the optional repeat keyword argument. For example,\n\
2518product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
Christian Heimesc3f30c42008-02-22 16:37:40 +00002519product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2520product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2521
2522static PyTypeObject product_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002523 PyVarObject_HEAD_INIT(NULL, 0)
2524 "itertools.product", /* tp_name */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07002525 sizeof(productobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002526 0, /* tp_itemsize */
2527 /* methods */
2528 (destructor)product_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002529 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002530 0, /* tp_getattr */
2531 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002532 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002533 0, /* tp_repr */
2534 0, /* tp_as_number */
2535 0, /* tp_as_sequence */
2536 0, /* tp_as_mapping */
2537 0, /* tp_hash */
2538 0, /* tp_call */
2539 0, /* tp_str */
2540 PyObject_GenericGetAttr, /* tp_getattro */
2541 0, /* tp_setattro */
2542 0, /* tp_as_buffer */
2543 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2544 Py_TPFLAGS_BASETYPE, /* tp_flags */
2545 product_doc, /* tp_doc */
2546 (traverseproc)product_traverse, /* tp_traverse */
2547 0, /* tp_clear */
2548 0, /* tp_richcompare */
2549 0, /* tp_weaklistoffset */
2550 PyObject_SelfIter, /* tp_iter */
2551 (iternextfunc)product_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002552 product_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002553 0, /* tp_members */
2554 0, /* tp_getset */
2555 0, /* tp_base */
2556 0, /* tp_dict */
2557 0, /* tp_descr_get */
2558 0, /* tp_descr_set */
2559 0, /* tp_dictoffset */
2560 0, /* tp_init */
2561 0, /* tp_alloc */
2562 product_new, /* tp_new */
2563 PyObject_GC_Del, /* tp_free */
Christian Heimesc3f30c42008-02-22 16:37:40 +00002564};
2565
2566
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002567/* combinations object *******************************************************/
Christian Heimes380f7f22008-02-28 11:19:05 +00002568
2569typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002570 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002571 PyObject *pool; /* input converted to a tuple */
2572 Py_ssize_t *indices; /* one index per result element */
2573 PyObject *result; /* most recently returned result tuple */
2574 Py_ssize_t r; /* size of result tuple */
2575 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimes380f7f22008-02-28 11:19:05 +00002576} combinationsobject;
2577
Tal Einatc4bccd32018-09-12 00:49:13 +03002578
2579/*[clinic input]
2580@classmethod
2581itertools.combinations.__new__
2582 iterable: object
2583 r: Py_ssize_t
2584Return successive r-length combinations of elements in the iterable.
2585
2586combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)
2587[clinic start generated code]*/
2588
Christian Heimes380f7f22008-02-28 11:19:05 +00002589static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002590itertools_combinations_impl(PyTypeObject *type, PyObject *iterable,
2591 Py_ssize_t r)
2592/*[clinic end generated code: output=87a689b39c40039c input=06bede09e3da20f8]*/
Christian Heimes380f7f22008-02-28 11:19:05 +00002593{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002594 combinationsobject *co;
2595 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002596 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002597 Py_ssize_t *indices = NULL;
2598 Py_ssize_t i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002599
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002600 pool = PySequence_Tuple(iterable);
2601 if (pool == NULL)
2602 goto error;
2603 n = PyTuple_GET_SIZE(pool);
2604 if (r < 0) {
2605 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2606 goto error;
2607 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002608
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002609 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002610 if (indices == NULL) {
2611 PyErr_NoMemory();
2612 goto error;
2613 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002614
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002615 for (i=0 ; i<r ; i++)
2616 indices[i] = i;
Christian Heimes380f7f22008-02-28 11:19:05 +00002617
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002618 /* create combinationsobject structure */
2619 co = (combinationsobject *)type->tp_alloc(type, 0);
2620 if (co == NULL)
2621 goto error;
Christian Heimes380f7f22008-02-28 11:19:05 +00002622
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002623 co->pool = pool;
2624 co->indices = indices;
2625 co->result = NULL;
2626 co->r = r;
2627 co->stopped = r > n ? 1 : 0;
2628
2629 return (PyObject *)co;
Christian Heimes380f7f22008-02-28 11:19:05 +00002630
2631error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002632 if (indices != NULL)
2633 PyMem_Free(indices);
2634 Py_XDECREF(pool);
2635 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002636}
2637
2638static void
2639combinations_dealloc(combinationsobject *co)
2640{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002641 PyObject_GC_UnTrack(co);
2642 Py_XDECREF(co->pool);
2643 Py_XDECREF(co->result);
2644 if (co->indices != NULL)
2645 PyMem_Free(co->indices);
2646 Py_TYPE(co)->tp_free(co);
Christian Heimes380f7f22008-02-28 11:19:05 +00002647}
2648
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002649static PyObject *
2650combinations_sizeof(combinationsobject *co, void *unused)
2651{
2652 Py_ssize_t res;
2653
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002654 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002655 res += co->r * sizeof(Py_ssize_t);
2656 return PyLong_FromSsize_t(res);
2657}
2658
Christian Heimes380f7f22008-02-28 11:19:05 +00002659static int
2660combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2661{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002662 Py_VISIT(co->pool);
2663 Py_VISIT(co->result);
2664 return 0;
Christian Heimes380f7f22008-02-28 11:19:05 +00002665}
2666
2667static PyObject *
2668combinations_next(combinationsobject *co)
2669{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002670 PyObject *elem;
2671 PyObject *oldelem;
2672 PyObject *pool = co->pool;
2673 Py_ssize_t *indices = co->indices;
2674 PyObject *result = co->result;
2675 Py_ssize_t n = PyTuple_GET_SIZE(pool);
2676 Py_ssize_t r = co->r;
2677 Py_ssize_t i, j, index;
Christian Heimes380f7f22008-02-28 11:19:05 +00002678
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002679 if (co->stopped)
2680 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002681
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002682 if (result == NULL) {
2683 /* On the first pass, initialize result tuple using the indices */
2684 result = PyTuple_New(r);
2685 if (result == NULL)
2686 goto empty;
2687 co->result = result;
2688 for (i=0; i<r ; i++) {
2689 index = indices[i];
2690 elem = PyTuple_GET_ITEM(pool, index);
2691 Py_INCREF(elem);
2692 PyTuple_SET_ITEM(result, i, elem);
2693 }
2694 } else {
2695 /* Copy the previous result tuple or re-use it if available */
2696 if (Py_REFCNT(result) > 1) {
2697 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05002698 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002699 if (result == NULL)
2700 goto empty;
2701 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002702 Py_DECREF(old_result);
2703 }
2704 /* Now, we've got the only copy so we can update it in-place
2705 * CPython's empty tuple is a singleton and cached in
2706 * PyTuple's freelist.
2707 */
2708 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimes380f7f22008-02-28 11:19:05 +00002709
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002710 /* Scan indices right-to-left until finding one that is not
2711 at its maximum (i + n - r). */
2712 for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2713 ;
Christian Heimes380f7f22008-02-28 11:19:05 +00002714
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002715 /* If i is negative, then the indices are all at
2716 their maximum value and we're done. */
2717 if (i < 0)
2718 goto empty;
Christian Heimes380f7f22008-02-28 11:19:05 +00002719
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002720 /* Increment the current index which we know is not at its
2721 maximum. Then move back to the right setting each index
2722 to its lowest possible value (one higher than the index
2723 to its left -- this maintains the sort order invariant). */
2724 indices[i]++;
2725 for (j=i+1 ; j<r ; j++)
2726 indices[j] = indices[j-1] + 1;
Christian Heimes380f7f22008-02-28 11:19:05 +00002727
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002728 /* Update the result tuple for the new indices
2729 starting with i, the leftmost index that changed */
2730 for ( ; i<r ; i++) {
2731 index = indices[i];
2732 elem = PyTuple_GET_ITEM(pool, index);
2733 Py_INCREF(elem);
2734 oldelem = PyTuple_GET_ITEM(result, i);
2735 PyTuple_SET_ITEM(result, i, elem);
2736 Py_DECREF(oldelem);
2737 }
2738 }
Christian Heimes380f7f22008-02-28 11:19:05 +00002739
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002740 Py_INCREF(result);
2741 return result;
Christian Heimes380f7f22008-02-28 11:19:05 +00002742
2743empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002744 co->stopped = 1;
2745 return NULL;
Christian Heimes380f7f22008-02-28 11:19:05 +00002746}
2747
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002748static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05302749combinations_reduce(combinationsobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002750{
2751 if (lz->result == NULL) {
2752 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2753 } else if (lz->stopped) {
2754 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2755 } else {
2756 PyObject *indices;
2757 Py_ssize_t i;
2758
2759 /* we must pickle the indices and use them for setstate */
2760 indices = PyTuple_New(lz->r);
2761 if (!indices)
2762 return NULL;
2763 for (i=0; i<lz->r; i++)
2764 {
2765 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2766 if (!index) {
2767 Py_DECREF(indices);
2768 return NULL;
2769 }
2770 PyTuple_SET_ITEM(indices, i, index);
2771 }
2772
2773 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2774 }
2775}
2776
2777static PyObject *
2778combinations_setstate(combinationsobject *lz, PyObject *state)
2779{
2780 PyObject *result;
2781 Py_ssize_t i;
2782 Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2783
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002784 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002785 PyErr_SetString(PyExc_ValueError, "invalid arguments");
2786 return NULL;
2787 }
2788
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002789 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002790 Py_ssize_t max;
2791 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2792 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002793
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002794 if (index == -1 && PyErr_Occurred())
2795 return NULL; /* not an integer */
2796 max = i + n - lz->r;
2797 /* clamp the index (beware of negative max) */
2798 if (index > max)
2799 index = max;
2800 if (index < 0)
2801 index = 0;
2802 lz->indices[i] = index;
2803 }
2804
2805 result = PyTuple_New(lz->r);
2806 if (result == NULL)
2807 return NULL;
2808 for (i=0; i<lz->r; i++) {
2809 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2810 Py_INCREF(element);
2811 PyTuple_SET_ITEM(result, i, element);
2812 }
2813
Serhiy Storchaka48842712016-04-06 09:45:48 +03002814 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002815 Py_RETURN_NONE;
2816}
2817
2818static PyMethodDef combinations_methods[] = {
2819 {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2820 reduce_doc},
2821 {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2822 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002823 {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2824 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002825 {NULL, NULL} /* sentinel */
2826};
2827
Christian Heimes380f7f22008-02-28 11:19:05 +00002828static PyTypeObject combinations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002829 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002830 "itertools.combinations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002831 sizeof(combinationsobject), /* tp_basicsize */
2832 0, /* tp_itemsize */
2833 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002834 (destructor)combinations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002835 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002836 0, /* tp_getattr */
2837 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02002838 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002839 0, /* tp_repr */
2840 0, /* tp_as_number */
2841 0, /* tp_as_sequence */
2842 0, /* tp_as_mapping */
2843 0, /* tp_hash */
2844 0, /* tp_call */
2845 0, /* tp_str */
2846 PyObject_GenericGetAttr, /* tp_getattro */
2847 0, /* tp_setattro */
2848 0, /* tp_as_buffer */
2849 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2850 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03002851 itertools_combinations__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002852 (traverseproc)combinations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002853 0, /* tp_clear */
2854 0, /* tp_richcompare */
2855 0, /* tp_weaklistoffset */
2856 PyObject_SelfIter, /* tp_iter */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00002857 (iternextfunc)combinations_next, /* tp_iternext */
2858 combinations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002859 0, /* tp_members */
2860 0, /* tp_getset */
2861 0, /* tp_base */
2862 0, /* tp_dict */
2863 0, /* tp_descr_get */
2864 0, /* tp_descr_set */
2865 0, /* tp_dictoffset */
2866 0, /* tp_init */
2867 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03002868 itertools_combinations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002869 PyObject_GC_Del, /* tp_free */
Christian Heimes380f7f22008-02-28 11:19:05 +00002870};
2871
2872
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002873/* combinations with replacement object **************************************/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002874
2875/* Equivalent to:
2876
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002877 def combinations_with_replacement(iterable, r):
2878 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2879 # number items returned: (n+r-1)! / r! / (n-1)!
2880 pool = tuple(iterable)
2881 n = len(pool)
2882 indices = [0] * r
2883 yield tuple(pool[i] for i in indices)
2884 while 1:
2885 for i in reversed(range(r)):
2886 if indices[i] != n - 1:
2887 break
2888 else:
2889 return
2890 indices[i:] = [indices[i] + 1] * (r - i)
2891 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002892
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002893 def combinations_with_replacement2(iterable, r):
2894 'Alternate version that filters from product()'
2895 pool = tuple(iterable)
2896 n = len(pool)
2897 for indices in product(range(n), repeat=r):
2898 if sorted(indices) == list(indices):
2899 yield tuple(pool[i] for i in indices)
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002900*/
2901typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002902 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07002903 PyObject *pool; /* input converted to a tuple */
2904 Py_ssize_t *indices; /* one index per result element */
2905 PyObject *result; /* most recently returned result tuple */
2906 Py_ssize_t r; /* size of result tuple */
2907 int stopped; /* set to 1 when the cwr iterator is exhausted */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002908} cwrobject;
2909
Tal Einatc4bccd32018-09-12 00:49:13 +03002910/*[clinic input]
2911@classmethod
2912itertools.combinations_with_replacement.__new__
2913 iterable: object
2914 r: Py_ssize_t
2915Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats.
2916
2917combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2918[clinic start generated code]*/
2919
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002920static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03002921itertools_combinations_with_replacement_impl(PyTypeObject *type,
2922 PyObject *iterable,
2923 Py_ssize_t r)
2924/*[clinic end generated code: output=48b26856d4e659ca input=dc2a8c7ba785fad7]*/
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002925{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002926 cwrobject *co;
2927 Py_ssize_t n;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002928 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002929 Py_ssize_t *indices = NULL;
2930 Py_ssize_t i;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002931
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002932 pool = PySequence_Tuple(iterable);
2933 if (pool == NULL)
2934 goto error;
2935 n = PyTuple_GET_SIZE(pool);
2936 if (r < 0) {
2937 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2938 goto error;
2939 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002940
Serhiy Storchakadee948b2015-02-03 01:34:09 +02002941 indices = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002942 if (indices == NULL) {
2943 PyErr_NoMemory();
2944 goto error;
2945 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002946
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002947 for (i=0 ; i<r ; i++)
2948 indices[i] = 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002949
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002950 /* create cwrobject structure */
2951 co = (cwrobject *)type->tp_alloc(type, 0);
2952 if (co == NULL)
2953 goto error;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002954
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002955 co->pool = pool;
2956 co->indices = indices;
2957 co->result = NULL;
2958 co->r = r;
2959 co->stopped = !n && r;
2960
2961 return (PyObject *)co;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002962
2963error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002964 if (indices != NULL)
2965 PyMem_Free(indices);
2966 Py_XDECREF(pool);
2967 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002968}
2969
2970static void
2971cwr_dealloc(cwrobject *co)
2972{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002973 PyObject_GC_UnTrack(co);
2974 Py_XDECREF(co->pool);
2975 Py_XDECREF(co->result);
2976 if (co->indices != NULL)
2977 PyMem_Free(co->indices);
2978 Py_TYPE(co)->tp_free(co);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002979}
2980
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002981static PyObject *
2982cwr_sizeof(cwrobject *co, void *unused)
2983{
2984 Py_ssize_t res;
2985
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02002986 res = _PyObject_SIZE(Py_TYPE(co));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02002987 res += co->r * sizeof(Py_ssize_t);
2988 return PyLong_FromSsize_t(res);
2989}
2990
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002991static int
2992cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2993{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00002994 Py_VISIT(co->pool);
2995 Py_VISIT(co->result);
2996 return 0;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00002997}
2998
2999static PyObject *
3000cwr_next(cwrobject *co)
3001{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003002 PyObject *elem;
3003 PyObject *oldelem;
3004 PyObject *pool = co->pool;
3005 Py_ssize_t *indices = co->indices;
3006 PyObject *result = co->result;
3007 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3008 Py_ssize_t r = co->r;
Tim Peters9edb1682013-09-03 11:49:31 -05003009 Py_ssize_t i, index;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003010
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003011 if (co->stopped)
3012 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003013
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003014 if (result == NULL) {
Tim Peters9edb1682013-09-03 11:49:31 -05003015 /* On the first pass, initialize result tuple with pool[0] */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003016 result = PyTuple_New(r);
3017 if (result == NULL)
3018 goto empty;
3019 co->result = result;
Raymond Hettingeracd61b62015-07-28 02:05:44 -07003020 if (n > 0) {
3021 elem = PyTuple_GET_ITEM(pool, 0);
3022 for (i=0; i<r ; i++) {
3023 assert(indices[i] == 0);
3024 Py_INCREF(elem);
3025 PyTuple_SET_ITEM(result, i, elem);
3026 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003027 }
3028 } else {
3029 /* Copy the previous result tuple or re-use it if available */
3030 if (Py_REFCNT(result) > 1) {
3031 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05003032 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003033 if (result == NULL)
3034 goto empty;
3035 co->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003036 Py_DECREF(old_result);
3037 }
3038 /* Now, we've got the only copy so we can update it in-place CPython's
3039 empty tuple is a singleton and cached in PyTuple's freelist. */
3040 assert(r == 0 || Py_REFCNT(result) == 1);
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003041
Tim Peters9edb1682013-09-03 11:49:31 -05003042 /* Scan indices right-to-left until finding one that is not
3043 * at its maximum (n-1). */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003044 for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
3045 ;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003046
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003047 /* If i is negative, then the indices are all at
Tim Peters9edb1682013-09-03 11:49:31 -05003048 their maximum value and we're done. */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003049 if (i < 0)
3050 goto empty;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003051
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003052 /* Increment the current index which we know is not at its
Tim Peters9edb1682013-09-03 11:49:31 -05003053 maximum. Then set all to the right to the same value. */
3054 index = indices[i] + 1;
3055 assert(index < n);
3056 elem = PyTuple_GET_ITEM(pool, index);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003057 for ( ; i<r ; i++) {
Tim Peters9edb1682013-09-03 11:49:31 -05003058 indices[i] = index;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003059 Py_INCREF(elem);
3060 oldelem = PyTuple_GET_ITEM(result, i);
3061 PyTuple_SET_ITEM(result, i, elem);
3062 Py_DECREF(oldelem);
3063 }
3064 }
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003065
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003066 Py_INCREF(result);
3067 return result;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003068
3069empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003070 co->stopped = 1;
3071 return NULL;
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003072}
3073
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003074static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303075cwr_reduce(cwrobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003076{
3077 if (lz->result == NULL) {
3078 return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
3079 } else if (lz->stopped) {
3080 return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
3081 } else {
3082 PyObject *indices;
3083 Py_ssize_t i;
3084
3085 /* we must pickle the indices and use them for setstate */
3086 indices = PyTuple_New(lz->r);
3087 if (!indices)
3088 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003089 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003090 PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
3091 if (!index) {
3092 Py_DECREF(indices);
3093 return NULL;
3094 }
3095 PyTuple_SET_ITEM(indices, i, index);
3096 }
3097
3098 return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
3099 }
3100}
3101
3102static PyObject *
3103cwr_setstate(cwrobject *lz, PyObject *state)
3104{
3105 PyObject *result;
3106 Py_ssize_t n, i;
3107
3108 if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
3109 {
3110 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3111 return NULL;
3112 }
3113
3114 n = PyTuple_GET_SIZE(lz->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003115 for (i=0; i<lz->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003116 PyObject* indexObject = PyTuple_GET_ITEM(state, i);
3117 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003118
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003119 if (index < 0 && PyErr_Occurred())
3120 return NULL; /* not an integer */
3121 /* clamp the index */
3122 if (index < 0)
3123 index = 0;
3124 else if (index > n-1)
3125 index = n-1;
3126 lz->indices[i] = index;
3127 }
3128 result = PyTuple_New(lz->r);
3129 if (result == NULL)
3130 return NULL;
3131 for (i=0; i<lz->r; i++) {
3132 PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
3133 Py_INCREF(element);
3134 PyTuple_SET_ITEM(result, i, element);
3135 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003136 Py_XSETREF(lz->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003137 Py_RETURN_NONE;
3138}
3139
3140static PyMethodDef cwr_methods[] = {
3141 {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
3142 reduce_doc},
3143 {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
3144 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003145 {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
3146 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003147 {NULL, NULL} /* sentinel */
3148};
3149
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003150static PyTypeObject cwr_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003151 PyVarObject_HEAD_INIT(NULL, 0)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003152 "itertools.combinations_with_replacement", /* tp_name */
3153 sizeof(cwrobject), /* tp_basicsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003154 0, /* tp_itemsize */
3155 /* methods */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003156 (destructor)cwr_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003157 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003158 0, /* tp_getattr */
3159 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003160 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003161 0, /* tp_repr */
3162 0, /* tp_as_number */
3163 0, /* tp_as_sequence */
3164 0, /* tp_as_mapping */
3165 0, /* tp_hash */
3166 0, /* tp_call */
3167 0, /* tp_str */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003168 PyObject_GenericGetAttr, /* tp_getattro */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003169 0, /* tp_setattro */
3170 0, /* tp_as_buffer */
3171 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003172 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003173 itertools_combinations_with_replacement__doc__, /* tp_doc */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003174 (traverseproc)cwr_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003175 0, /* tp_clear */
3176 0, /* tp_richcompare */
3177 0, /* tp_weaklistoffset */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003178 PyObject_SelfIter, /* tp_iter */
3179 (iternextfunc)cwr_next, /* tp_iternext */
3180 cwr_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003181 0, /* tp_members */
3182 0, /* tp_getset */
3183 0, /* tp_base */
3184 0, /* tp_dict */
3185 0, /* tp_descr_get */
3186 0, /* tp_descr_set */
3187 0, /* tp_dictoffset */
3188 0, /* tp_init */
3189 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003190 itertools_combinations_with_replacement, /* tp_new */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003191 PyObject_GC_Del, /* tp_free */
Raymond Hettingerd07d9392009-01-27 04:20:44 +00003192};
3193
3194
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003195/* permutations object ********************************************************
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003196
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003197def permutations(iterable, r=None):
Sergey0078a0c2019-10-29 08:10:24 +03003198 # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
3199 # permutations(range(3)) --> 012 021 102 120 201 210
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003200 pool = tuple(iterable)
3201 n = len(pool)
3202 r = n if r is None else r
Sergey0078a0c2019-10-29 08:10:24 +03003203 if r > n:
3204 return
3205 indices = list(range(n))
3206 cycles = list(range(n, n-r, -1))
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003207 yield tuple(pool[i] for i in indices[:r])
3208 while n:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003209 for i in reversed(range(r)):
3210 cycles[i] -= 1
3211 if cycles[i] == 0:
3212 indices[i:] = indices[i+1:] + indices[i:i+1]
3213 cycles[i] = n - i
3214 else:
3215 j = cycles[i]
3216 indices[i], indices[-j] = indices[-j], indices[i]
3217 yield tuple(pool[i] for i in indices[:r])
3218 break
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003219 else:
Serhiy Storchakad741a882015-06-11 00:06:39 +03003220 return
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003221*/
3222
3223typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003224 PyObject_HEAD
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003225 PyObject *pool; /* input converted to a tuple */
3226 Py_ssize_t *indices; /* one index per element in the pool */
3227 Py_ssize_t *cycles; /* one rollover counter per element in the result */
3228 PyObject *result; /* most recently returned result tuple */
3229 Py_ssize_t r; /* size of result tuple */
3230 int stopped; /* set to 1 when the iterator is exhausted */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003231} permutationsobject;
3232
Tal Einatc4bccd32018-09-12 00:49:13 +03003233/*[clinic input]
3234@classmethod
3235itertools.permutations.__new__
3236 iterable: object
3237 r as robj: object = None
3238Return successive r-length permutations of elements in the iterable.
3239
3240permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)
3241[clinic start generated code]*/
3242
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003243static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003244itertools_permutations_impl(PyTypeObject *type, PyObject *iterable,
3245 PyObject *robj)
3246/*[clinic end generated code: output=296a72fa76d620ea input=57d0170a4ac0ec7a]*/
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003247{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003248 permutationsobject *po;
3249 Py_ssize_t n;
3250 Py_ssize_t r;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003251 PyObject *pool = NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003252 Py_ssize_t *indices = NULL;
3253 Py_ssize_t *cycles = NULL;
3254 Py_ssize_t i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003255
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003256 pool = PySequence_Tuple(iterable);
3257 if (pool == NULL)
3258 goto error;
3259 n = PyTuple_GET_SIZE(pool);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003260
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003261 r = n;
3262 if (robj != Py_None) {
3263 if (!PyLong_Check(robj)) {
3264 PyErr_SetString(PyExc_TypeError, "Expected int as r");
3265 goto error;
3266 }
3267 r = PyLong_AsSsize_t(robj);
3268 if (r == -1 && PyErr_Occurred())
3269 goto error;
3270 }
3271 if (r < 0) {
3272 PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3273 goto error;
3274 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003275
Serhiy Storchakadee948b2015-02-03 01:34:09 +02003276 indices = PyMem_New(Py_ssize_t, n);
3277 cycles = PyMem_New(Py_ssize_t, r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003278 if (indices == NULL || cycles == NULL) {
3279 PyErr_NoMemory();
3280 goto error;
3281 }
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003282
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003283 for (i=0 ; i<n ; i++)
3284 indices[i] = i;
3285 for (i=0 ; i<r ; i++)
3286 cycles[i] = n - i;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003287
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003288 /* create permutationsobject structure */
3289 po = (permutationsobject *)type->tp_alloc(type, 0);
3290 if (po == NULL)
3291 goto error;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003292
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003293 po->pool = pool;
3294 po->indices = indices;
3295 po->cycles = cycles;
3296 po->result = NULL;
3297 po->r = r;
3298 po->stopped = r > n ? 1 : 0;
3299
3300 return (PyObject *)po;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003301
3302error:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003303 if (indices != NULL)
3304 PyMem_Free(indices);
3305 if (cycles != NULL)
3306 PyMem_Free(cycles);
3307 Py_XDECREF(pool);
3308 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003309}
3310
3311static void
3312permutations_dealloc(permutationsobject *po)
3313{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003314 PyObject_GC_UnTrack(po);
3315 Py_XDECREF(po->pool);
3316 Py_XDECREF(po->result);
3317 PyMem_Free(po->indices);
3318 PyMem_Free(po->cycles);
3319 Py_TYPE(po)->tp_free(po);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003320}
3321
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003322static PyObject *
3323permutations_sizeof(permutationsobject *po, void *unused)
3324{
3325 Py_ssize_t res;
3326
Serhiy Storchaka5c4064e2015-12-19 20:05:25 +02003327 res = _PyObject_SIZE(Py_TYPE(po));
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003328 res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3329 res += po->r * sizeof(Py_ssize_t);
3330 return PyLong_FromSsize_t(res);
3331}
3332
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003333static int
3334permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3335{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003336 Py_VISIT(po->pool);
3337 Py_VISIT(po->result);
3338 return 0;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003339}
3340
3341static PyObject *
3342permutations_next(permutationsobject *po)
3343{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003344 PyObject *elem;
3345 PyObject *oldelem;
3346 PyObject *pool = po->pool;
3347 Py_ssize_t *indices = po->indices;
3348 Py_ssize_t *cycles = po->cycles;
3349 PyObject *result = po->result;
3350 Py_ssize_t n = PyTuple_GET_SIZE(pool);
3351 Py_ssize_t r = po->r;
3352 Py_ssize_t i, j, k, index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003353
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003354 if (po->stopped)
3355 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003356
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003357 if (result == NULL) {
3358 /* On the first pass, initialize result tuple using the indices */
3359 result = PyTuple_New(r);
3360 if (result == NULL)
3361 goto empty;
3362 po->result = result;
3363 for (i=0; i<r ; i++) {
3364 index = indices[i];
3365 elem = PyTuple_GET_ITEM(pool, index);
3366 Py_INCREF(elem);
3367 PyTuple_SET_ITEM(result, i, elem);
3368 }
3369 } else {
3370 if (n == 0)
3371 goto empty;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003372
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003373 /* Copy the previous result tuple or re-use it if available */
3374 if (Py_REFCNT(result) > 1) {
3375 PyObject *old_result = result;
Sergey Fedoseev234531b2019-02-25 21:59:12 +05003376 result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003377 if (result == NULL)
3378 goto empty;
3379 po->result = result;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003380 Py_DECREF(old_result);
3381 }
3382 /* Now, we've got the only copy so we can update it in-place */
3383 assert(r == 0 || Py_REFCNT(result) == 1);
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003384
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003385 /* Decrement rightmost cycle, moving leftward upon zero rollover */
3386 for (i=r-1 ; i>=0 ; i--) {
3387 cycles[i] -= 1;
3388 if (cycles[i] == 0) {
3389 /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3390 index = indices[i];
3391 for (j=i ; j<n-1 ; j++)
3392 indices[j] = indices[j+1];
3393 indices[n-1] = index;
3394 cycles[i] = n - i;
3395 } else {
3396 j = cycles[i];
3397 index = indices[i];
3398 indices[i] = indices[n-j];
3399 indices[n-j] = index;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003400
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003401 for (k=i; k<r ; k++) {
3402 /* start with i, the leftmost element that changed */
3403 /* yield tuple(pool[k] for k in indices[:r]) */
3404 index = indices[k];
3405 elem = PyTuple_GET_ITEM(pool, index);
3406 Py_INCREF(elem);
3407 oldelem = PyTuple_GET_ITEM(result, k);
3408 PyTuple_SET_ITEM(result, k, elem);
3409 Py_DECREF(oldelem);
3410 }
3411 break;
3412 }
3413 }
3414 /* If i is negative, then the cycles have all
3415 rolled-over and we're done. */
3416 if (i < 0)
3417 goto empty;
3418 }
3419 Py_INCREF(result);
3420 return result;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003421
3422empty:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003423 po->stopped = 1;
3424 return NULL;
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003425}
3426
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003427static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303428permutations_reduce(permutationsobject *po, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003429{
3430 if (po->result == NULL) {
3431 return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3432 } else if (po->stopped) {
3433 return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3434 } else {
3435 PyObject *indices=NULL, *cycles=NULL;
3436 Py_ssize_t n, i;
3437
3438 /* we must pickle the indices and cycles and use them for setstate */
3439 n = PyTuple_GET_SIZE(po->pool);
3440 indices = PyTuple_New(n);
3441 if (indices == NULL)
3442 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003443 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003444 PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3445 if (!index)
3446 goto err;
3447 PyTuple_SET_ITEM(indices, i, index);
3448 }
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003449
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003450 cycles = PyTuple_New(po->r);
3451 if (cycles == NULL)
3452 goto err;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003453 for (i=0 ; i<po->r ; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003454 PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3455 if (!index)
3456 goto err;
3457 PyTuple_SET_ITEM(cycles, i, index);
3458 }
3459 return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3460 po->pool, po->r,
3461 indices, cycles);
3462 err:
3463 Py_XDECREF(indices);
3464 Py_XDECREF(cycles);
3465 return NULL;
3466 }
3467}
3468
3469static PyObject *
3470permutations_setstate(permutationsobject *po, PyObject *state)
3471{
3472 PyObject *indices, *cycles, *result;
3473 Py_ssize_t n, i;
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003474
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003475 if (!PyTuple_Check(state)) {
3476 PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3477 return NULL;
3478 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003479 if (!PyArg_ParseTuple(state, "O!O!",
3480 &PyTuple_Type, &indices,
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003481 &PyTuple_Type, &cycles)) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003482 return NULL;
Serhiy Storchaka85c3f262016-10-02 08:34:53 +03003483 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003484
3485 n = PyTuple_GET_SIZE(po->pool);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003486 if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003487 PyErr_SetString(PyExc_ValueError, "invalid arguments");
3488 return NULL;
3489 }
3490
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003491 for (i=0; i<n; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003492 PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3493 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3494 if (index < 0 && PyErr_Occurred())
3495 return NULL; /* not an integer */
3496 /* clamp the index */
3497 if (index < 0)
3498 index = 0;
3499 else if (index > n-1)
3500 index = n-1;
3501 po->indices[i] = index;
3502 }
3503
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003504 for (i=0; i<po->r; i++) {
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003505 PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3506 Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3507 if (index < 0 && PyErr_Occurred())
3508 return NULL; /* not an integer */
3509 if (index < 1)
3510 index = 1;
3511 else if (index > n-i)
3512 index = n-i;
3513 po->cycles[i] = index;
3514 }
3515 result = PyTuple_New(po->r);
3516 if (result == NULL)
3517 return NULL;
3518 for (i=0; i<po->r; i++) {
3519 PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3520 Py_INCREF(element);
3521 PyTuple_SET_ITEM(result, i, element);
3522 }
Serhiy Storchaka48842712016-04-06 09:45:48 +03003523 Py_XSETREF(po->result, result);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003524 Py_RETURN_NONE;
3525}
3526
3527static PyMethodDef permuations_methods[] = {
3528 {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3529 reduce_doc},
3530 {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3531 setstate_doc},
Serhiy Storchaka2dae92a2013-12-09 17:45:57 +02003532 {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3533 sizeof_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003534 {NULL, NULL} /* sentinel */
3535};
3536
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003537static PyTypeObject permutations_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003538 PyVarObject_HEAD_INIT(NULL, 0)
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003539 "itertools.permutations", /* tp_name */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003540 sizeof(permutationsobject), /* tp_basicsize */
3541 0, /* tp_itemsize */
3542 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003543 (destructor)permutations_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003544 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003545 0, /* tp_getattr */
3546 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003547 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003548 0, /* tp_repr */
3549 0, /* tp_as_number */
3550 0, /* tp_as_sequence */
3551 0, /* tp_as_mapping */
3552 0, /* tp_hash */
3553 0, /* tp_call */
3554 0, /* tp_str */
3555 PyObject_GenericGetAttr, /* tp_getattro */
3556 0, /* tp_setattro */
3557 0, /* tp_as_buffer */
3558 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3559 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003560 itertools_permutations__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003561 (traverseproc)permutations_traverse,/* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003562 0, /* tp_clear */
3563 0, /* tp_richcompare */
3564 0, /* tp_weaklistoffset */
3565 PyObject_SelfIter, /* tp_iter */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003566 (iternextfunc)permutations_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003567 permuations_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003568 0, /* tp_members */
3569 0, /* tp_getset */
3570 0, /* tp_base */
3571 0, /* tp_dict */
3572 0, /* tp_descr_get */
3573 0, /* tp_descr_set */
3574 0, /* tp_dictoffset */
3575 0, /* tp_init */
3576 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003577 itertools_permutations, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003578 PyObject_GC_Del, /* tp_free */
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003579};
3580
Tal Einatc4bccd32018-09-12 00:49:13 +03003581
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003582/* accumulate object ********************************************************/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003583
3584typedef struct {
3585 PyObject_HEAD
3586 PyObject *total;
3587 PyObject *it;
Raymond Hettinger5d446132011-03-27 18:52:10 -07003588 PyObject *binop;
Lisa Roach9718b592018-09-23 17:34:59 -07003589 PyObject *initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003590} accumulateobject;
3591
Tal Einatc4bccd32018-09-12 00:49:13 +03003592/*[clinic input]
3593@classmethod
3594itertools.accumulate.__new__
3595 iterable: object
3596 func as binop: object = None
Lisa Roach9718b592018-09-23 17:34:59 -07003597 *
3598 initial: object = None
Tal Einatc4bccd32018-09-12 00:49:13 +03003599Return series of accumulated sums (or other binary function results).
3600[clinic start generated code]*/
3601
Raymond Hettinger482ba772010-12-01 22:48:00 +00003602static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003603itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
Lisa Roach9718b592018-09-23 17:34:59 -07003604 PyObject *binop, PyObject *initial)
3605/*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
Raymond Hettinger482ba772010-12-01 22:48:00 +00003606{
Raymond Hettinger482ba772010-12-01 22:48:00 +00003607 PyObject *it;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003608 accumulateobject *lz;
3609
Raymond Hettinger482ba772010-12-01 22:48:00 +00003610 /* Get iterator. */
3611 it = PyObject_GetIter(iterable);
3612 if (it == NULL)
3613 return NULL;
3614
Raymond Hettinger482ba772010-12-01 22:48:00 +00003615 /* create accumulateobject structure */
3616 lz = (accumulateobject *)type->tp_alloc(type, 0);
3617 if (lz == NULL) {
3618 Py_DECREF(it);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003619 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003620 }
3621
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003622 if (binop != Py_None) {
3623 Py_XINCREF(binop);
3624 lz->binop = binop;
3625 }
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003626 lz->total = NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003627 lz->it = it;
Lisa Roach9718b592018-09-23 17:34:59 -07003628 Py_XINCREF(initial);
3629 lz->initial = initial;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003630 return (PyObject *)lz;
3631}
3632
3633static void
3634accumulate_dealloc(accumulateobject *lz)
3635{
3636 PyObject_GC_UnTrack(lz);
Raymond Hettinger5d446132011-03-27 18:52:10 -07003637 Py_XDECREF(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003638 Py_XDECREF(lz->total);
3639 Py_XDECREF(lz->it);
Lisa Roach9718b592018-09-23 17:34:59 -07003640 Py_XDECREF(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003641 Py_TYPE(lz)->tp_free(lz);
3642}
3643
3644static int
3645accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3646{
Raymond Hettinger5d446132011-03-27 18:52:10 -07003647 Py_VISIT(lz->binop);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003648 Py_VISIT(lz->it);
3649 Py_VISIT(lz->total);
Lisa Roach9718b592018-09-23 17:34:59 -07003650 Py_VISIT(lz->initial);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003651 return 0;
3652}
3653
3654static PyObject *
3655accumulate_next(accumulateobject *lz)
3656{
Serhiy Storchaka1ed017a2015-12-27 15:51:32 +02003657 PyObject *val, *newtotal;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003658
Lisa Roach9718b592018-09-23 17:34:59 -07003659 if (lz->initial != Py_None) {
3660 lz->total = lz->initial;
3661 Py_INCREF(Py_None);
3662 lz->initial = Py_None;
3663 Py_INCREF(lz->total);
3664 return lz->total;
3665 }
Raymond Hettingerb5244a32015-08-16 14:24:20 -07003666 val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003667 if (val == NULL)
3668 return NULL;
Serhiy Storchakae7e9c322013-01-25 13:37:39 +02003669
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003670 if (lz->total == NULL) {
3671 Py_INCREF(val);
3672 lz->total = val;
3673 return lz->total;
3674 }
Raymond Hettinger5d446132011-03-27 18:52:10 -07003675
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02003676 if (lz->binop == NULL)
Raymond Hettinger5d446132011-03-27 18:52:10 -07003677 newtotal = PyNumber_Add(lz->total, val);
3678 else
3679 newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003680 Py_DECREF(val);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003681 if (newtotal == NULL)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +00003682 return NULL;
Raymond Hettinger482ba772010-12-01 22:48:00 +00003683
Raymond Hettinger482ba772010-12-01 22:48:00 +00003684 Py_INCREF(newtotal);
Serhiy Storchakaf01e4082016-04-10 18:12:01 +03003685 Py_SETREF(lz->total, newtotal);
Raymond Hettinger482ba772010-12-01 22:48:00 +00003686 return newtotal;
3687}
3688
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003689static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303690accumulate_reduce(accumulateobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003691{
Lisa Roach9718b592018-09-23 17:34:59 -07003692 if (lz->initial != Py_None) {
3693 PyObject *it;
3694
3695 assert(lz->total == NULL);
3696 if (PyType_Ready(&chain_type) < 0)
3697 return NULL;
3698 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3699 lz->initial, lz->it);
3700 if (it == NULL)
3701 return NULL;
3702 return Py_BuildValue("O(NO)O", Py_TYPE(lz),
3703 it, lz->binop?lz->binop:Py_None, Py_None);
3704 }
Serhiy Storchakad5516252016-03-06 14:00:45 +02003705 if (lz->total == Py_None) {
3706 PyObject *it;
3707
3708 if (PyType_Ready(&chain_type) < 0)
3709 return NULL;
3710 if (PyType_Ready(&islice_type) < 0)
3711 return NULL;
3712 it = PyObject_CallFunction((PyObject *)&chain_type, "(O)O",
3713 lz->total, lz->it);
3714 if (it == NULL)
3715 return NULL;
3716 it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3717 it, lz->binop ? lz->binop : Py_None);
3718 if (it == NULL)
3719 return NULL;
3720 return Py_BuildValue("O(NiO)", &islice_type, it, 1, Py_None);
3721 }
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003722 return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3723 lz->it, lz->binop?lz->binop:Py_None,
3724 lz->total?lz->total:Py_None);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003725}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003726
3727static PyObject *
3728accumulate_setstate(accumulateobject *lz, PyObject *state)
3729{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02003730 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03003731 Py_XSETREF(lz->total, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003732 Py_RETURN_NONE;
3733}
3734
3735static PyMethodDef accumulate_methods[] = {
3736 {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3737 reduce_doc},
3738 {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3739 setstate_doc},
3740 {NULL, NULL} /* sentinel */
3741};
3742
Raymond Hettinger482ba772010-12-01 22:48:00 +00003743static PyTypeObject accumulate_type = {
3744 PyVarObject_HEAD_INIT(NULL, 0)
3745 "itertools.accumulate", /* tp_name */
3746 sizeof(accumulateobject), /* tp_basicsize */
3747 0, /* tp_itemsize */
3748 /* methods */
3749 (destructor)accumulate_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003750 0, /* tp_vectorcall_offset */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003751 0, /* tp_getattr */
3752 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003753 0, /* tp_as_async */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003754 0, /* tp_repr */
3755 0, /* tp_as_number */
3756 0, /* tp_as_sequence */
3757 0, /* tp_as_mapping */
3758 0, /* tp_hash */
3759 0, /* tp_call */
3760 0, /* tp_str */
3761 PyObject_GenericGetAttr, /* tp_getattro */
3762 0, /* tp_setattro */
3763 0, /* tp_as_buffer */
3764 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
3765 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003766 itertools_accumulate__doc__, /* tp_doc */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003767 (traverseproc)accumulate_traverse, /* tp_traverse */
3768 0, /* tp_clear */
3769 0, /* tp_richcompare */
3770 0, /* tp_weaklistoffset */
3771 PyObject_SelfIter, /* tp_iter */
3772 (iternextfunc)accumulate_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003773 accumulate_methods, /* tp_methods */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003774 0, /* tp_members */
3775 0, /* tp_getset */
3776 0, /* tp_base */
3777 0, /* tp_dict */
3778 0, /* tp_descr_get */
3779 0, /* tp_descr_set */
3780 0, /* tp_dictoffset */
3781 0, /* tp_init */
3782 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003783 itertools_accumulate, /* tp_new */
Raymond Hettinger482ba772010-12-01 22:48:00 +00003784 PyObject_GC_Del, /* tp_free */
3785};
3786
Christian Heimesdd15f6c2008-03-16 00:07:10 +00003787
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003788/* compress object ************************************************************/
3789
3790/* Equivalent to:
3791
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003792 def compress(data, selectors):
3793 "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3794 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003795*/
3796
3797typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003798 PyObject_HEAD
3799 PyObject *data;
3800 PyObject *selectors;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003801} compressobject;
3802
Tal Einatc4bccd32018-09-12 00:49:13 +03003803/*[clinic input]
3804@classmethod
3805itertools.compress.__new__
3806 data as seq1: object
3807 selectors as seq2: object
3808Return data elements corresponding to true selector elements.
3809
3810Forms a shorter iterator from selected data elements using the selectors to
3811choose the data elements.
3812[clinic start generated code]*/
3813
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003814static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003815itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
3816/*[clinic end generated code: output=7e67157212ed09e0 input=79596d7cd20c77e5]*/
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003817{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003818 PyObject *data=NULL, *selectors=NULL;
3819 compressobject *lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003820
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003821 data = PyObject_GetIter(seq1);
3822 if (data == NULL)
3823 goto fail;
3824 selectors = PyObject_GetIter(seq2);
3825 if (selectors == NULL)
3826 goto fail;
3827
3828 /* create compressobject structure */
3829 lz = (compressobject *)type->tp_alloc(type, 0);
3830 if (lz == NULL)
3831 goto fail;
3832 lz->data = data;
3833 lz->selectors = selectors;
3834 return (PyObject *)lz;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003835
3836fail:
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003837 Py_XDECREF(data);
3838 Py_XDECREF(selectors);
3839 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003840}
3841
3842static void
3843compress_dealloc(compressobject *lz)
3844{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003845 PyObject_GC_UnTrack(lz);
3846 Py_XDECREF(lz->data);
3847 Py_XDECREF(lz->selectors);
3848 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003849}
3850
3851static int
3852compress_traverse(compressobject *lz, visitproc visit, void *arg)
3853{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003854 Py_VISIT(lz->data);
3855 Py_VISIT(lz->selectors);
3856 return 0;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003857}
3858
3859static PyObject *
3860compress_next(compressobject *lz)
3861{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003862 PyObject *data = lz->data, *selectors = lz->selectors;
3863 PyObject *datum, *selector;
3864 PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3865 PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3866 int ok;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003867
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003868 while (1) {
3869 /* Steps: get datum, get selector, evaluate selector.
3870 Order is important (to match the pure python version
3871 in terms of which input gets a chance to raise an
3872 exception first).
3873 */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003874
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003875 datum = datanext(data);
Serhiy Storchaka278d03b2013-04-06 22:52:34 +03003876 if (datum == NULL)
Serhiy Storchakae8f706e2013-04-06 21:14:43 +03003877 return NULL;
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003878
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003879 selector = selectornext(selectors);
3880 if (selector == NULL) {
3881 Py_DECREF(datum);
3882 return NULL;
3883 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003884
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003885 ok = PyObject_IsTrue(selector);
3886 Py_DECREF(selector);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003887 if (ok > 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003888 return datum;
3889 Py_DECREF(datum);
Antoine Pitrou721738f2012-08-15 23:20:39 +02003890 if (ok < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003891 return NULL;
3892 }
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003893}
3894
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003895static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05303896compress_reduce(compressobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003897{
3898 return Py_BuildValue("O(OO)", Py_TYPE(lz),
3899 lz->data, lz->selectors);
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07003900}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00003901
3902static PyMethodDef compress_methods[] = {
3903 {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3904 reduce_doc},
3905 {NULL, NULL} /* sentinel */
3906};
3907
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003908static PyTypeObject compress_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003909 PyVarObject_HEAD_INIT(NULL, 0)
3910 "itertools.compress", /* tp_name */
3911 sizeof(compressobject), /* tp_basicsize */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003912 0, /* tp_itemsize */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003913 /* methods */
3914 (destructor)compress_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003915 0, /* tp_vectorcall_offset */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003916 0, /* tp_getattr */
3917 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02003918 0, /* tp_as_async */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003919 0, /* tp_repr */
3920 0, /* tp_as_number */
3921 0, /* tp_as_sequence */
3922 0, /* tp_as_mapping */
3923 0, /* tp_hash */
3924 0, /* tp_call */
3925 0, /* tp_str */
3926 PyObject_GenericGetAttr, /* tp_getattro */
3927 0, /* tp_setattro */
3928 0, /* tp_as_buffer */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003929 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003930 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03003931 itertools_compress__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003932 (traverseproc)compress_traverse, /* tp_traverse */
3933 0, /* tp_clear */
3934 0, /* tp_richcompare */
3935 0, /* tp_weaklistoffset */
3936 PyObject_SelfIter, /* tp_iter */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003937 (iternextfunc)compress_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003938 compress_methods, /* tp_methods */
3939 0, /* tp_members */
3940 0, /* tp_getset */
3941 0, /* tp_base */
3942 0, /* tp_dict */
3943 0, /* tp_descr_get */
3944 0, /* tp_descr_set */
3945 0, /* tp_dictoffset */
3946 0, /* tp_init */
3947 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03003948 itertools_compress, /* tp_new */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07003949 PyObject_GC_Del, /* tp_free */
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +00003950};
3951
3952
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003953/* filterfalse object ************************************************************/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003954
3955typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003956 PyObject_HEAD
3957 PyObject *func;
3958 PyObject *it;
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003959} filterfalseobject;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003960
Tal Einatc4bccd32018-09-12 00:49:13 +03003961/*[clinic input]
3962@classmethod
3963itertools.filterfalse.__new__
3964 function as func: object
3965 iterable as seq: object
3966 /
3967Return those items of iterable for which function(item) is false.
3968
3969If function is None, return the items that are false.
3970[clinic start generated code]*/
3971
Raymond Hettinger60eca932003-02-09 06:40:58 +00003972static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03003973itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
3974/*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
Raymond Hettinger60eca932003-02-09 06:40:58 +00003975{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003976 PyObject *it;
3977 filterfalseobject *lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003978
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003979 /* Get iterator. */
3980 it = PyObject_GetIter(seq);
3981 if (it == NULL)
3982 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003983
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003984 /* create filterfalseobject structure */
3985 lz = (filterfalseobject *)type->tp_alloc(type, 0);
3986 if (lz == NULL) {
3987 Py_DECREF(it);
3988 return NULL;
3989 }
3990 Py_INCREF(func);
3991 lz->func = func;
3992 lz->it = it;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003993
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00003994 return (PyObject *)lz;
Raymond Hettinger60eca932003-02-09 06:40:58 +00003995}
3996
3997static void
Raymond Hettingerb0002d22008-03-13 01:41:43 +00003998filterfalse_dealloc(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00003999{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004000 PyObject_GC_UnTrack(lz);
4001 Py_XDECREF(lz->func);
4002 Py_XDECREF(lz->it);
4003 Py_TYPE(lz)->tp_free(lz);
Raymond Hettinger60eca932003-02-09 06:40:58 +00004004}
4005
4006static int
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004007filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
Raymond Hettinger60eca932003-02-09 06:40:58 +00004008{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004009 Py_VISIT(lz->it);
4010 Py_VISIT(lz->func);
4011 return 0;
Raymond Hettinger60eca932003-02-09 06:40:58 +00004012}
4013
4014static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004015filterfalse_next(filterfalseobject *lz)
Raymond Hettinger60eca932003-02-09 06:40:58 +00004016{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004017 PyObject *item;
4018 PyObject *it = lz->it;
4019 long ok;
4020 PyObject *(*iternext)(PyObject *);
Raymond Hettinger60eca932003-02-09 06:40:58 +00004021
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004022 iternext = *Py_TYPE(it)->tp_iternext;
4023 for (;;) {
4024 item = iternext(it);
4025 if (item == NULL)
4026 return NULL;
Raymond Hettinger60eca932003-02-09 06:40:58 +00004027
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004028 if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
4029 ok = PyObject_IsTrue(item);
4030 } else {
4031 PyObject *good;
Petr Viktorinffd97532020-02-11 17:46:57 +01004032 good = PyObject_CallOneArg(lz->func, item);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004033 if (good == NULL) {
4034 Py_DECREF(item);
4035 return NULL;
4036 }
4037 ok = PyObject_IsTrue(good);
4038 Py_DECREF(good);
4039 }
Antoine Pitrou6f430e42012-08-15 23:18:25 +02004040 if (ok == 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004041 return item;
4042 Py_DECREF(item);
Antoine Pitrou6f430e42012-08-15 23:18:25 +02004043 if (ok < 0)
4044 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004045 }
Raymond Hettinger60eca932003-02-09 06:40:58 +00004046}
4047
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004048static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304049filterfalse_reduce(filterfalseobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004050{
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004051 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
4052}
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004053
4054static PyMethodDef filterfalse_methods[] = {
4055 {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
4056 reduce_doc},
4057 {NULL, NULL} /* sentinel */
4058};
4059
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004060static PyTypeObject filterfalse_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004061 PyVarObject_HEAD_INIT(NULL, 0)
4062 "itertools.filterfalse", /* tp_name */
4063 sizeof(filterfalseobject), /* tp_basicsize */
4064 0, /* tp_itemsize */
4065 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004066 (destructor)filterfalse_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004067 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004068 0, /* tp_getattr */
4069 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004070 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004071 0, /* tp_repr */
4072 0, /* tp_as_number */
4073 0, /* tp_as_sequence */
4074 0, /* tp_as_mapping */
4075 0, /* tp_hash */
4076 0, /* tp_call */
4077 0, /* tp_str */
4078 PyObject_GenericGetAttr, /* tp_getattro */
4079 0, /* tp_setattro */
4080 0, /* tp_as_buffer */
4081 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4082 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03004083 itertools_filterfalse__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004084 (traverseproc)filterfalse_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004085 0, /* tp_clear */
4086 0, /* tp_richcompare */
4087 0, /* tp_weaklistoffset */
4088 PyObject_SelfIter, /* tp_iter */
4089 (iternextfunc)filterfalse_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004090 filterfalse_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004091 0, /* tp_members */
4092 0, /* tp_getset */
4093 0, /* tp_base */
4094 0, /* tp_dict */
4095 0, /* tp_descr_get */
4096 0, /* tp_descr_set */
4097 0, /* tp_dictoffset */
4098 0, /* tp_init */
4099 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03004100 itertools_filterfalse, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004101 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004102};
4103
4104
4105/* count object ************************************************************/
4106
4107typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004108 PyObject_HEAD
4109 Py_ssize_t cnt;
4110 PyObject *long_cnt;
4111 PyObject *long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004112} countobject;
4113
Raymond Hettinger30729212009-02-12 06:28:27 +00004114/* Counting logic and invariants:
4115
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004116fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
Raymond Hettinger30729212009-02-12 06:28:27 +00004117
Serhiy Storchaka483405b2015-02-17 10:14:30 +02004118 assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004119 Advances with: cnt += 1
4120 When count hits Y_SSIZE_T_MAX, switch to slow_mode.
Raymond Hettinger30729212009-02-12 06:28:27 +00004121
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004122slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
Raymond Hettinger30729212009-02-12 06:28:27 +00004123
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004124 assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
4125 All counting is done with python objects (no overflows or underflows).
4126 Advances with: long_cnt += long_step
4127 Step may be zero -- effectively a slow version of repeat(cnt).
4128 Either long_cnt or long_step may be a float, Fraction, or Decimal.
Raymond Hettinger30729212009-02-12 06:28:27 +00004129*/
4130
Tal Einatc4bccd32018-09-12 00:49:13 +03004131/*[clinic input]
4132@classmethod
4133itertools.count.__new__
4134 start as long_cnt: object(c_default="NULL") = 0
4135 step as long_step: object(c_default="NULL") = 1
4136Return a count object whose .__next__() method returns consecutive values.
4137
4138Equivalent to:
4139 def count(firstval=0, step=1):
4140 x = firstval
4141 while 1:
4142 yield x
4143 x += step
4144[clinic start generated code]*/
4145
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004146static PyObject *
Tal Einatc4bccd32018-09-12 00:49:13 +03004147itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
4148 PyObject *long_step)
4149/*[clinic end generated code: output=09a9250aebd00b1c input=d7a85eec18bfcd94]*/
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004150{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004151 countobject *lz;
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004152 int fast_mode;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004153 Py_ssize_t cnt = 0;
Kristjan Valur Jonsson35722a92011-03-30 11:04:28 +00004154 long step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004155
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004156 if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
4157 (long_step != NULL && !PyNumber_Check(long_step))) {
4158 PyErr_SetString(PyExc_TypeError, "a number is required");
4159 return NULL;
4160 }
Raymond Hettinger30729212009-02-12 06:28:27 +00004161
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004162 fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
4163 (long_step == NULL || PyLong_Check(long_step));
4164
4165 /* If not specified, start defaults to 0 */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004166 if (long_cnt != NULL) {
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004167 if (fast_mode) {
4168 assert(PyLong_Check(long_cnt));
4169 cnt = PyLong_AsSsize_t(long_cnt);
4170 if (cnt == -1 && PyErr_Occurred()) {
4171 PyErr_Clear();
4172 fast_mode = 0;
4173 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004174 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004175 } else {
4176 cnt = 0;
Victor Stinner37834132020-10-27 17:12:53 +01004177 long_cnt = _PyLong_GetZero();
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004178 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004179 Py_INCREF(long_cnt);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004180
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004181 /* If not specified, step defaults to 1 */
Victor Stinner37834132020-10-27 17:12:53 +01004182 if (long_step == NULL) {
4183 long_step = _PyLong_GetOne();
4184 }
Serhiy Storchakaba85d692017-03-30 09:09:41 +03004185 Py_INCREF(long_step);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004186
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004187 assert(long_cnt != NULL && long_step != NULL);
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004188
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004189 /* Fast mode only works when the step is 1 */
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004190 if (fast_mode) {
4191 assert(PyLong_Check(long_step));
4192 step = PyLong_AsLong(long_step);
4193 if (step != 1) {
4194 fast_mode = 0;
4195 if (step == -1 && PyErr_Occurred())
4196 PyErr_Clear();
4197 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004198 }
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004199
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004200 if (fast_mode)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004201 Py_CLEAR(long_cnt);
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004202 else
4203 cnt = PY_SSIZE_T_MAX;
Raymond Hettingereb13fdd2009-02-21 22:30:12 +00004204
Serhiy Storchaka8ddcf3a2016-09-10 09:49:24 +03004205 assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4206 (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4207 assert(!fast_mode ||
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004208 (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004209
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004210 /* create countobject structure */
4211 lz = (countobject *)type->tp_alloc(type, 0);
4212 if (lz == NULL) {
4213 Py_XDECREF(long_cnt);
Zackery Spytz0523c392019-03-26 00:05:29 -06004214 Py_DECREF(long_step);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004215 return NULL;
4216 }
4217 lz->cnt = cnt;
4218 lz->long_cnt = long_cnt;
4219 lz->long_step = long_step;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004220
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004221 return (PyObject *)lz;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004222}
4223
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004224static void
4225count_dealloc(countobject *lz)
4226{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004227 PyObject_GC_UnTrack(lz);
4228 Py_XDECREF(lz->long_cnt);
4229 Py_XDECREF(lz->long_step);
4230 Py_TYPE(lz)->tp_free(lz);
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004231}
4232
Benjamin Peterson25e26a02009-02-16 21:28:29 +00004233static int
Raymond Hettingerd6280f42009-02-16 20:50:56 +00004234count_traverse(countobject *lz, visitproc visit, void *arg)
4235{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004236 Py_VISIT(lz->long_cnt);
4237 Py_VISIT(lz->long_step);
4238 return 0;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004239}
4240
4241static PyObject *
4242count_nextlong(countobject *lz)
4243{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004244 PyObject *long_cnt;
4245 PyObject *stepped_up;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004246
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004247 long_cnt = lz->long_cnt;
4248 if (long_cnt == NULL) {
4249 /* Switch to slow_mode */
4250 long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4251 if (long_cnt == NULL)
4252 return NULL;
4253 }
4254 assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
Raymond Hettinger30729212009-02-12 06:28:27 +00004255
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004256 stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4257 if (stepped_up == NULL)
4258 return NULL;
4259 lz->long_cnt = stepped_up;
4260 return long_cnt;
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004261}
4262
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004263static PyObject *
4264count_next(countobject *lz)
4265{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004266 if (lz->cnt == PY_SSIZE_T_MAX)
4267 return count_nextlong(lz);
4268 return PyLong_FromSsize_t(lz->cnt++);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004269}
4270
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004271static PyObject *
4272count_repr(countobject *lz)
4273{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004274 if (lz->cnt != PY_SSIZE_T_MAX)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004275 return PyUnicode_FromFormat("%s(%zd)",
4276 _PyType_Name(Py_TYPE(lz)), lz->cnt);
Guido van Rossum8ce8a782007-11-01 19:42:39 +00004277
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004278 if (PyLong_Check(lz->long_step)) {
4279 long step = PyLong_AsLong(lz->long_step);
4280 if (step == -1 && PyErr_Occurred()) {
4281 PyErr_Clear();
4282 }
4283 if (step == 1) {
4284 /* Don't display step when it is an integer equal to 1 */
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004285 return PyUnicode_FromFormat("%s(%R)",
4286 _PyType_Name(Py_TYPE(lz)),
4287 lz->long_cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004288 }
4289 }
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004290 return PyUnicode_FromFormat("%s(%R, %R)",
4291 _PyType_Name(Py_TYPE(lz)),
4292 lz->long_cnt, lz->long_step);
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004293}
4294
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004295static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304296count_reduce(countobject *lz, PyObject *Py_UNUSED(ignored))
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004297{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004298 if (lz->cnt == PY_SSIZE_T_MAX)
4299 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4300 return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004301}
4302
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004303static PyMethodDef count_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004304 {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004305 reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004306 {NULL, NULL} /* sentinel */
Raymond Hettingera3e1ad22009-11-30 22:02:31 +00004307};
4308
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004309static PyTypeObject count_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004310 PyVarObject_HEAD_INIT(NULL, 0)
4311 "itertools.count", /* tp_name */
4312 sizeof(countobject), /* tp_basicsize */
4313 0, /* tp_itemsize */
4314 /* methods */
4315 (destructor)count_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004316 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004317 0, /* tp_getattr */
4318 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004319 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004320 (reprfunc)count_repr, /* tp_repr */
4321 0, /* tp_as_number */
4322 0, /* tp_as_sequence */
4323 0, /* tp_as_mapping */
4324 0, /* tp_hash */
4325 0, /* tp_call */
4326 0, /* tp_str */
4327 PyObject_GenericGetAttr, /* tp_getattro */
4328 0, /* tp_setattro */
4329 0, /* tp_as_buffer */
4330 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004331 Py_TPFLAGS_BASETYPE, /* tp_flags */
Tal Einatc4bccd32018-09-12 00:49:13 +03004332 itertools_count__doc__, /* tp_doc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004333 (traverseproc)count_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004334 0, /* tp_clear */
4335 0, /* tp_richcompare */
4336 0, /* tp_weaklistoffset */
4337 PyObject_SelfIter, /* tp_iter */
4338 (iternextfunc)count_next, /* tp_iternext */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004339 count_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004340 0, /* tp_members */
4341 0, /* tp_getset */
4342 0, /* tp_base */
4343 0, /* tp_dict */
4344 0, /* tp_descr_get */
4345 0, /* tp_descr_set */
4346 0, /* tp_dictoffset */
4347 0, /* tp_init */
4348 0, /* tp_alloc */
Tal Einatc4bccd32018-09-12 00:49:13 +03004349 itertools_count, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004350 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004351};
4352
4353
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004354/* repeat object ************************************************************/
4355
4356typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004357 PyObject_HEAD
4358 PyObject *element;
4359 Py_ssize_t cnt;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004360} repeatobject;
4361
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004362static PyTypeObject repeat_type;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004363
4364static PyObject *
4365repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4366{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004367 repeatobject *ro;
4368 PyObject *element;
Raymond Hettingeraf464502019-11-09 20:28:31 -08004369 Py_ssize_t cnt = -1, n_args;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004370 static char *kwargs[] = {"object", "times", NULL};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004371
Raymond Hettingeraf464502019-11-09 20:28:31 -08004372 n_args = PyTuple_GET_SIZE(args);
4373 if (kwds != NULL)
4374 n_args += PyDict_GET_SIZE(kwds);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004375 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4376 &element, &cnt))
4377 return NULL;
Raymond Hettinger97d35552014-06-24 21:36:58 -07004378 /* Does user supply times argument? */
Raymond Hettingeraf464502019-11-09 20:28:31 -08004379 if (n_args == 2 && cnt < 0)
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004380 cnt = 0;
4381
4382 ro = (repeatobject *)type->tp_alloc(type, 0);
4383 if (ro == NULL)
4384 return NULL;
4385 Py_INCREF(element);
4386 ro->element = element;
4387 ro->cnt = cnt;
4388 return (PyObject *)ro;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004389}
4390
4391static void
4392repeat_dealloc(repeatobject *ro)
4393{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004394 PyObject_GC_UnTrack(ro);
4395 Py_XDECREF(ro->element);
4396 Py_TYPE(ro)->tp_free(ro);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004397}
4398
4399static int
4400repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4401{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004402 Py_VISIT(ro->element);
4403 return 0;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004404}
4405
4406static PyObject *
4407repeat_next(repeatobject *ro)
4408{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004409 if (ro->cnt == 0)
4410 return NULL;
4411 if (ro->cnt > 0)
4412 ro->cnt--;
4413 Py_INCREF(ro->element);
4414 return ro->element;
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004415}
4416
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004417static PyObject *
4418repeat_repr(repeatobject *ro)
4419{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004420 if (ro->cnt == -1)
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004421 return PyUnicode_FromFormat("%s(%R)",
4422 _PyType_Name(Py_TYPE(ro)), ro->element);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004423 else
Serhiy Storchakab3a77962017-09-21 14:24:13 +03004424 return PyUnicode_FromFormat("%s(%R, %zd)",
4425 _PyType_Name(Py_TYPE(ro)), ro->element,
4426 ro->cnt);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004427}
Raymond Hettinger7dacda22004-04-08 21:54:00 +00004428
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004429static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304430repeat_len(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004431{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004432 if (ro->cnt == -1) {
4433 PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4434 return NULL;
4435 }
4436 return PyLong_FromSize_t(ro->cnt);
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004437}
4438
Armin Rigof5b3e362006-02-11 21:32:43 +00004439PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004440
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004441static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304442repeat_reduce(repeatobject *ro, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004443{
4444 /* unpickle this so that a new repeat iterator is constructed with an
4445 * object, then call __setstate__ on it to set cnt
4446 */
4447 if (ro->cnt >= 0)
4448 return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4449 else
4450 return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4451}
4452
Raymond Hettinger6b27cda2005-09-24 21:23:05 +00004453static PyMethodDef repeat_methods[] = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004454 {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004455 {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004456 {NULL, NULL} /* sentinel */
Raymond Hettinger5cab2e32004-02-10 09:25:40 +00004457};
4458
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004459PyDoc_STRVAR(repeat_doc,
Raymond Hettingerf4bb7f22009-02-19 02:44:01 +00004460"repeat(object [,times]) -> create an iterator which returns the object\n\
4461for the specified number of times. If not specified, returns the object\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004462endlessly.");
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004463
Raymond Hettinger1d7a3482003-07-14 07:07:12 +00004464static PyTypeObject repeat_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004465 PyVarObject_HEAD_INIT(NULL, 0)
4466 "itertools.repeat", /* tp_name */
4467 sizeof(repeatobject), /* tp_basicsize */
4468 0, /* tp_itemsize */
4469 /* methods */
4470 (destructor)repeat_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004471 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004472 0, /* tp_getattr */
4473 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004474 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004475 (reprfunc)repeat_repr, /* tp_repr */
4476 0, /* tp_as_number */
4477 0, /* tp_as_sequence */
4478 0, /* tp_as_mapping */
4479 0, /* tp_hash */
4480 0, /* tp_call */
4481 0, /* tp_str */
4482 PyObject_GenericGetAttr, /* tp_getattro */
4483 0, /* tp_setattro */
4484 0, /* tp_as_buffer */
4485 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4486 Py_TPFLAGS_BASETYPE, /* tp_flags */
4487 repeat_doc, /* tp_doc */
4488 (traverseproc)repeat_traverse, /* tp_traverse */
4489 0, /* tp_clear */
4490 0, /* tp_richcompare */
4491 0, /* tp_weaklistoffset */
4492 PyObject_SelfIter, /* tp_iter */
4493 (iternextfunc)repeat_next, /* tp_iternext */
4494 repeat_methods, /* tp_methods */
4495 0, /* tp_members */
4496 0, /* tp_getset */
4497 0, /* tp_base */
4498 0, /* tp_dict */
4499 0, /* tp_descr_get */
4500 0, /* tp_descr_set */
4501 0, /* tp_dictoffset */
4502 0, /* tp_init */
4503 0, /* tp_alloc */
4504 repeat_new, /* tp_new */
4505 PyObject_GC_Del, /* tp_free */
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004506};
4507
Tal Einatc4bccd32018-09-12 00:49:13 +03004508
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004509/* ziplongest object *********************************************************/
Thomas Wouterscf297e42007-02-23 15:07:44 +00004510
4511typedef struct {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004512 PyObject_HEAD
4513 Py_ssize_t tuplesize;
4514 Py_ssize_t numactive;
4515 PyObject *ittuple; /* tuple of iterators */
4516 PyObject *result;
4517 PyObject *fillvalue;
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004518} ziplongestobject;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004519
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004520static PyTypeObject ziplongest_type;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004521
4522static PyObject *
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004523zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004524{
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004525 _Py_IDENTIFIER(fillvalue);
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004526 ziplongestobject *lz;
4527 Py_ssize_t i;
4528 PyObject *ittuple; /* tuple of iterators */
4529 PyObject *result;
4530 PyObject *fillvalue = Py_None;
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004531 Py_ssize_t tuplesize;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004532
Serhiy Storchaka5ab81d72016-12-16 16:18:57 +02004533 if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
Serhiy Storchakaa24107b2019-02-25 17:59:46 +02004534 fillvalue = NULL;
4535 if (PyDict_GET_SIZE(kwds) == 1) {
4536 fillvalue = _PyDict_GetItemIdWithError(kwds, &PyId_fillvalue);
4537 }
4538 if (fillvalue == NULL) {
4539 if (!PyErr_Occurred()) {
4540 PyErr_SetString(PyExc_TypeError,
4541 "zip_longest() got an unexpected keyword argument");
4542 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004543 return NULL;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004544 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004545 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004546
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004547 /* args must be a tuple */
4548 assert(PyTuple_Check(args));
Serhiy Storchakabf623ae2017-04-19 20:03:52 +03004549 tuplesize = PyTuple_GET_SIZE(args);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004550
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004551 /* obtain iterators */
4552 ittuple = PyTuple_New(tuplesize);
4553 if (ittuple == NULL)
4554 return NULL;
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004555 for (i=0; i < tuplesize; i++) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004556 PyObject *item = PyTuple_GET_ITEM(args, i);
4557 PyObject *it = PyObject_GetIter(item);
4558 if (it == NULL) {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004559 Py_DECREF(ittuple);
4560 return NULL;
4561 }
4562 PyTuple_SET_ITEM(ittuple, i, it);
4563 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004564
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004565 /* create a result holder */
4566 result = PyTuple_New(tuplesize);
4567 if (result == NULL) {
4568 Py_DECREF(ittuple);
4569 return NULL;
4570 }
4571 for (i=0 ; i < tuplesize ; i++) {
4572 Py_INCREF(Py_None);
4573 PyTuple_SET_ITEM(result, i, Py_None);
4574 }
Thomas Wouterscf297e42007-02-23 15:07:44 +00004575
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004576 /* create ziplongestobject structure */
4577 lz = (ziplongestobject *)type->tp_alloc(type, 0);
4578 if (lz == NULL) {
4579 Py_DECREF(ittuple);
4580 Py_DECREF(result);
4581 return NULL;
4582 }
4583 lz->ittuple = ittuple;
4584 lz->tuplesize = tuplesize;
4585 lz->numactive = tuplesize;
4586 lz->result = result;
4587 Py_INCREF(fillvalue);
4588 lz->fillvalue = fillvalue;
4589 return (PyObject *)lz;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004590}
4591
4592static void
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004593zip_longest_dealloc(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004594{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004595 PyObject_GC_UnTrack(lz);
4596 Py_XDECREF(lz->ittuple);
4597 Py_XDECREF(lz->result);
4598 Py_XDECREF(lz->fillvalue);
4599 Py_TYPE(lz)->tp_free(lz);
Thomas Wouterscf297e42007-02-23 15:07:44 +00004600}
4601
4602static int
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004603zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004604{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004605 Py_VISIT(lz->ittuple);
4606 Py_VISIT(lz->result);
4607 Py_VISIT(lz->fillvalue);
4608 return 0;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004609}
4610
4611static PyObject *
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004612zip_longest_next(ziplongestobject *lz)
Thomas Wouterscf297e42007-02-23 15:07:44 +00004613{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004614 Py_ssize_t i;
4615 Py_ssize_t tuplesize = lz->tuplesize;
4616 PyObject *result = lz->result;
4617 PyObject *it;
4618 PyObject *item;
4619 PyObject *olditem;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004620
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004621 if (tuplesize == 0)
4622 return NULL;
4623 if (lz->numactive == 0)
4624 return NULL;
4625 if (Py_REFCNT(result) == 1) {
4626 Py_INCREF(result);
4627 for (i=0 ; i < tuplesize ; i++) {
4628 it = PyTuple_GET_ITEM(lz->ittuple, i);
4629 if (it == NULL) {
4630 Py_INCREF(lz->fillvalue);
4631 item = lz->fillvalue;
4632 } else {
4633 item = PyIter_Next(it);
4634 if (item == NULL) {
4635 lz->numactive -= 1;
4636 if (lz->numactive == 0 || PyErr_Occurred()) {
4637 lz->numactive = 0;
4638 Py_DECREF(result);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004639 return NULL;
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004640 } else {
4641 Py_INCREF(lz->fillvalue);
4642 item = lz->fillvalue;
4643 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4644 Py_DECREF(it);
4645 }
Raymond Hettingerfc438512009-11-01 20:55:33 +00004646 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004647 }
4648 olditem = PyTuple_GET_ITEM(result, i);
4649 PyTuple_SET_ITEM(result, i, item);
4650 Py_DECREF(olditem);
Raymond Hettingerfc438512009-11-01 20:55:33 +00004651 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004652 } else {
4653 result = PyTuple_New(tuplesize);
4654 if (result == NULL)
4655 return NULL;
4656 for (i=0 ; i < tuplesize ; i++) {
4657 it = PyTuple_GET_ITEM(lz->ittuple, i);
4658 if (it == NULL) {
4659 Py_INCREF(lz->fillvalue);
4660 item = lz->fillvalue;
4661 } else {
4662 item = PyIter_Next(it);
4663 if (item == NULL) {
4664 lz->numactive -= 1;
4665 if (lz->numactive == 0 || PyErr_Occurred()) {
4666 lz->numactive = 0;
4667 Py_DECREF(result);
4668 return NULL;
4669 } else {
4670 Py_INCREF(lz->fillvalue);
4671 item = lz->fillvalue;
4672 PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4673 Py_DECREF(it);
4674 }
4675 }
4676 }
4677 PyTuple_SET_ITEM(result, i, item);
4678 }
4679 }
4680 return result;
Thomas Wouterscf297e42007-02-23 15:07:44 +00004681}
4682
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004683static PyObject *
Siddhesh Poyarekar55edd0c2018-04-30 00:29:33 +05304684zip_longest_reduce(ziplongestobject *lz, PyObject *Py_UNUSED(ignored))
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004685{
Serhiy Storchakad269b5e2013-01-25 13:38:56 +02004686
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004687 /* Create a new tuple with empty sequences where appropriate to pickle.
4688 * Then use setstate to set the fillvalue
4689 */
4690 int i;
4691 PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
Raymond Hettingera6ea44a2015-08-17 23:55:28 -07004692
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004693 if (args == NULL)
4694 return NULL;
4695 for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4696 PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4697 if (elem == NULL) {
4698 elem = PyTuple_New(0);
4699 if (elem == NULL) {
4700 Py_DECREF(args);
4701 return NULL;
4702 }
4703 } else
4704 Py_INCREF(elem);
4705 PyTuple_SET_ITEM(args, i, elem);
4706 }
4707 return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4708}
4709
4710static PyObject *
4711zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4712{
Serhiy Storchaka4a1e70f2015-12-27 12:36:18 +02004713 Py_INCREF(state);
Serhiy Storchaka48842712016-04-06 09:45:48 +03004714 Py_XSETREF(lz->fillvalue, state);
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004715 Py_RETURN_NONE;
4716}
4717
4718static PyMethodDef zip_longest_methods[] = {
4719 {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4720 reduce_doc},
4721 {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4722 setstate_doc},
4723 {NULL, NULL} /* sentinel */
4724};
4725
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004726PyDoc_STRVAR(zip_longest_doc,
4727"zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004728\n\
Serhiy Storchaka6a7b3a72016-04-17 08:32:47 +03004729Return a zip_longest object whose .__next__() method returns a tuple where\n\
Georg Brandla18af4e2007-04-21 15:47:16 +00004730the i-th element comes from the i-th iterable argument. The .__next__()\n\
Thomas Wouterscf297e42007-02-23 15:07:44 +00004731method continues until the longest iterable in the argument sequence\n\
4732is exhausted and then it raises StopIteration. When the shorter iterables\n\
4733are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4734defaults to None or can be specified by a keyword argument.\n\
4735");
4736
Raymond Hettinger736c0ab2008-03-13 02:09:15 +00004737static PyTypeObject ziplongest_type = {
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004738 PyVarObject_HEAD_INIT(NULL, 0)
4739 "itertools.zip_longest", /* tp_name */
4740 sizeof(ziplongestobject), /* tp_basicsize */
4741 0, /* tp_itemsize */
4742 /* methods */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004743 (destructor)zip_longest_dealloc, /* tp_dealloc */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004744 0, /* tp_vectorcall_offset */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004745 0, /* tp_getattr */
4746 0, /* tp_setattr */
Jeroen Demeyer530f5062019-05-31 04:13:39 +02004747 0, /* tp_as_async */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004748 0, /* tp_repr */
4749 0, /* tp_as_number */
4750 0, /* tp_as_sequence */
4751 0, /* tp_as_mapping */
4752 0, /* tp_hash */
4753 0, /* tp_call */
4754 0, /* tp_str */
4755 PyObject_GenericGetAttr, /* tp_getattro */
4756 0, /* tp_setattro */
4757 0, /* tp_as_buffer */
4758 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
4759 Py_TPFLAGS_BASETYPE, /* tp_flags */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004760 zip_longest_doc, /* tp_doc */
4761 (traverseproc)zip_longest_traverse, /* tp_traverse */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004762 0, /* tp_clear */
4763 0, /* tp_richcompare */
4764 0, /* tp_weaklistoffset */
4765 PyObject_SelfIter, /* tp_iter */
4766 (iternextfunc)zip_longest_next, /* tp_iternext */
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004767 zip_longest_methods, /* tp_methods */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004768 0, /* tp_members */
4769 0, /* tp_getset */
4770 0, /* tp_base */
4771 0, /* tp_dict */
4772 0, /* tp_descr_get */
4773 0, /* tp_descr_set */
4774 0, /* tp_dictoffset */
4775 0, /* tp_init */
4776 0, /* tp_alloc */
Raymond Hettingerb468e1f2015-08-14 14:10:49 -07004777 zip_longest_new, /* tp_new */
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004778 PyObject_GC_Del, /* tp_free */
Thomas Wouterscf297e42007-02-23 15:07:44 +00004779};
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004780
Tal Einatc4bccd32018-09-12 00:49:13 +03004781
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004782/* module level code ********************************************************/
4783
4784PyDoc_STRVAR(module_doc,
4785"Functional tools for creating and using iterators.\n\
4786\n\
4787Infinite iterators:\n\
Andrew Kuchlingb003ffa2013-06-21 07:58:35 -04004788count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
Raymond Hettinger61fe64d2003-02-23 04:40:07 +00004789cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
Andrew M. Kuchlingdff694b2003-04-14 15:31:27 +00004790repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004791\n\
4792Iterators terminating on the shortest input sequence:\n\
Raymond Hettinger5bf70912011-03-27 18:59:51 -07004793accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
oldkaa0735f2018-02-02 16:52:55 +08004794chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
4795chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
Raymond Hettinger3522a582009-01-29 03:41:55 +00004796compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4797dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4798groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
Raymond Hettingerb0002d22008-03-13 01:41:43 +00004799filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004800islice(seq, [start,] stop [, step]) --> elements from\n\
4801 seq[start:stop:step]\n\
Raymond Hettingercc061d02020-11-30 20:42:54 -08004802pairwise(s) --> (s[0],s[1]), (s[1],s[2]), (s[2], s[3]), ...\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004803starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
Raymond Hettingerad983e72003-11-12 14:32:26 +00004804tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004805takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
oldkaa0735f2018-02-02 16:52:55 +08004806zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
Raymond Hettinger811f3dc2009-01-29 03:48:02 +00004807\n\
4808Combinatoric generators:\n\
4809product(p, q, ... [repeat=1]) --> cartesian product\n\
4810permutations(p[, r])\n\
Raymond Hettinger36c3c022009-11-19 01:20:07 +00004811combinations(p, r)\n\
4812combinations_with_replacement(p, r)\n\
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004813");
4814
Dong-hee Na514c4692020-03-18 02:46:24 +09004815static int
4816itertoolsmodule_exec(PyObject *m)
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004817{
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004818 PyTypeObject *typelist[] = {
Raymond Hettinger482ba772010-12-01 22:48:00 +00004819 &accumulate_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004820 &combinations_type,
4821 &cwr_type,
4822 &cycle_type,
4823 &dropwhile_type,
4824 &takewhile_type,
4825 &islice_type,
4826 &starmap_type,
4827 &chain_type,
4828 &compress_type,
4829 &filterfalse_type,
4830 &count_type,
4831 &ziplongest_type,
Raymond Hettingercc061d02020-11-30 20:42:54 -08004832 &pairwise_type,
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004833 &permutations_type,
4834 &product_type,
4835 &repeat_type,
4836 &groupby_type,
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00004837 &_grouper_type,
4838 &tee_type,
Dong-hee Na05e4a292020-03-23 01:17:34 +09004839 &teedataobject_type
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004840 };
Raymond Hettinger60eca932003-02-09 06:40:58 +00004841
Victor Stinnerd2ec81a2020-02-07 09:17:07 +01004842 Py_SET_TYPE(&teedataobject_type, &PyType_Type);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004843
Dong-hee Na05e4a292020-03-23 01:17:34 +09004844 for (size_t i = 0; i < Py_ARRAY_LENGTH(typelist); i++) {
4845 if (PyModule_AddType(m, typelist[i]) < 0) {
Dong-hee Na514c4692020-03-18 02:46:24 +09004846 return -1;
4847 }
Antoine Pitrouf95a1b32010-05-09 15:52:27 +00004848 }
Raymond Hettingerad983e72003-11-12 14:32:26 +00004849
Dong-hee Na514c4692020-03-18 02:46:24 +09004850 return 0;
4851}
4852
4853static struct PyModuleDef_Slot itertoolsmodule_slots[] = {
4854 {Py_mod_exec, itertoolsmodule_exec},
4855 {0, NULL}
4856};
4857
4858static PyMethodDef module_methods[] = {
4859 ITERTOOLS_TEE_METHODDEF
4860 {NULL, NULL} /* sentinel */
4861};
4862
4863
4864static struct PyModuleDef itertoolsmodule = {
4865 PyModuleDef_HEAD_INIT,
4866 "itertools",
4867 module_doc,
4868 0,
4869 module_methods,
4870 itertoolsmodule_slots,
4871 NULL,
4872 NULL,
4873 NULL
4874};
4875
4876PyMODINIT_FUNC
4877PyInit_itertools(void)
4878{
4879 return PyModuleDef_Init(&itertoolsmodule);
Raymond Hettinger96ef8112003-02-01 00:10:11 +00004880}