Branch data Line data Source code
1 : : #define PY_SSIZE_T_CLEAN
2 : : #include "Python.h"
3 : : #include "pycore_call.h" // _PyObject_CallNoArgs()
4 : : #include "pycore_long.h" // _PyLong_GetZero()
5 : : #include "pycore_moduleobject.h" // _PyModule_GetState()
6 : : #include "pycore_typeobject.h" // _PyType_GetModuleState()
7 : : #include "pycore_object.h" // _PyObject_GC_TRACK()
8 : : #include "pycore_tuple.h" // _PyTuple_ITEMS()
9 : : #include "structmember.h" // PyMemberDef
10 : : #include <stddef.h> // offsetof()
11 : :
12 : : /* Itertools module written and maintained
13 : : by Raymond D. Hettinger <python@rcn.com>
14 : : */
15 : :
16 : : typedef struct {
17 : : PyTypeObject *accumulate_type;
18 : : PyTypeObject *batched_type;
19 : : PyTypeObject *chain_type;
20 : : PyTypeObject *combinations_type;
21 : : PyTypeObject *compress_type;
22 : : PyTypeObject *count_type;
23 : : PyTypeObject *cwr_type;
24 : : PyTypeObject *cycle_type;
25 : : PyTypeObject *dropwhile_type;
26 : : PyTypeObject *filterfalse_type;
27 : : PyTypeObject *groupby_type;
28 : : PyTypeObject *_grouper_type;
29 : : PyTypeObject *islice_type;
30 : : PyTypeObject *pairwise_type;
31 : : PyTypeObject *permutations_type;
32 : : PyTypeObject *product_type;
33 : : PyTypeObject *repeat_type;
34 : : PyTypeObject *starmap_type;
35 : : PyTypeObject *takewhile_type;
36 : : PyTypeObject *tee_type;
37 : : PyTypeObject *teedataobject_type;
38 : : PyTypeObject *ziplongest_type;
39 : : } itertools_state;
40 : :
41 : : static inline itertools_state *
42 : 122 : get_module_state(PyObject *mod)
43 : : {
44 : 122 : void *state = _PyModule_GetState(mod);
45 : : assert(state != NULL);
46 : 122 : return (itertools_state *)state;
47 : : }
48 : :
49 : : static inline itertools_state *
50 : 0 : get_module_state_by_cls(PyTypeObject *cls)
51 : : {
52 : 0 : void *state = _PyType_GetModuleState(cls);
53 : : assert(state != NULL);
54 : 0 : return (itertools_state *)state;
55 : : }
56 : :
57 : : static struct PyModuleDef itertoolsmodule;
58 : :
59 : : static inline itertools_state *
60 : 18 : find_state_by_type(PyTypeObject *tp)
61 : : {
62 : 18 : PyObject *mod = PyType_GetModuleByDef(tp, &itertoolsmodule);
63 : : assert(mod != NULL);
64 : 18 : return get_module_state(mod);
65 : : }
66 : :
67 : : /*[clinic input]
68 : : module itertools
69 : : class itertools.groupby "groupbyobject *" "clinic_state()->groupby_type"
70 : : class itertools._grouper "_grouperobject *" "clinic_state()->_grouper_type"
71 : : class itertools.teedataobject "teedataobject *" "clinic_state()->teedataobject_type"
72 : : class itertools._tee "teeobject *" "clinic_state()->tee_type"
73 : : class itertools.batched "batchedobject *" "clinic_state()->batched_type"
74 : : class itertools.cycle "cycleobject *" "clinic_state()->cycle_type"
75 : : class itertools.dropwhile "dropwhileobject *" "clinic_state()->dropwhile_type"
76 : : class itertools.takewhile "takewhileobject *" "clinic_state()->takewhile_type"
77 : : class itertools.starmap "starmapobject *" "clinic_state()->starmap_type"
78 : : class itertools.chain "chainobject *" "clinic_state()->chain_type"
79 : : class itertools.combinations "combinationsobject *" "clinic_state()->combinations_type"
80 : : class itertools.combinations_with_replacement "cwr_object *" "clinic_state()->cwr_type"
81 : : class itertools.permutations "permutationsobject *" "clinic_state()->permutations_type"
82 : : class itertools.accumulate "accumulateobject *" "clinic_state()->accumulate_type"
83 : : class itertools.compress "compressobject *" "clinic_state()->compress_type"
84 : : class itertools.filterfalse "filterfalseobject *" "clinic_state()->filterfalse_type"
85 : : class itertools.count "countobject *" "clinic_state()->count_type"
86 : : class itertools.pairwise "pairwiseobject *" "clinic_state()->pairwise_type"
87 : : [clinic start generated code]*/
88 : : /*[clinic end generated code: output=da39a3ee5e6b4b0d input=aa48fe4de9d4080f]*/
89 : :
90 : : #define clinic_state() (find_state_by_type(type))
91 : : #define clinic_state_by_cls() (get_module_state_by_cls(base_tp))
92 : : #include "clinic/itertoolsmodule.c.h"
93 : : #undef clinic_state_by_cls
94 : : #undef clinic_state
95 : :
96 : : /* batched object ************************************************************/
97 : :
98 : : /* Note: The built-in zip() function includes a "strict" argument
99 : : that was needed because that function would silently truncate data,
100 : : and there was no easy way for a user to detect the data loss.
101 : : The same reasoning does not apply to batched() which never drops data.
102 : : Instead, batched() produces a shorter tuple which can be handled
103 : : as the user sees fit. If requested, it would be reasonable to add
104 : : "fillvalue" support which had demonstrated value in zip_longest().
105 : : For now, the API is kept simple and clean.
106 : : */
107 : :
108 : : typedef struct {
109 : : PyObject_HEAD
110 : : PyObject *it;
111 : : Py_ssize_t batch_size;
112 : : } batchedobject;
113 : :
114 : : /*[clinic input]
115 : : @classmethod
116 : : itertools.batched.__new__ as batched_new
117 : : iterable: object
118 : : n: Py_ssize_t
119 : : Batch data into tuples of length n. The last batch may be shorter than n.
120 : :
121 : : Loops over the input iterable and accumulates data into tuples
122 : : up to size n. The input is consumed lazily, just enough to
123 : : fill a batch. The result is yielded as soon as a batch is full
124 : : or when the input iterable is exhausted.
125 : :
126 : : >>> for batch in batched('ABCDEFG', 3):
127 : : ... print(batch)
128 : : ...
129 : : ('A', 'B', 'C')
130 : : ('D', 'E', 'F')
131 : : ('G',)
132 : :
133 : : [clinic start generated code]*/
134 : :
135 : : static PyObject *
136 : 0 : batched_new_impl(PyTypeObject *type, PyObject *iterable, Py_ssize_t n)
137 : : /*[clinic end generated code: output=7ebc954d655371b6 input=ffd70726927c5129]*/
138 : : {
139 : : PyObject *it;
140 : : batchedobject *bo;
141 : :
142 [ # # ]: 0 : if (n < 1) {
143 : : /* We could define the n==0 case to return an empty iterator
144 : : but that is at odds with the idea that batching should
145 : : never throw-away input data.
146 : : */
147 : 0 : PyErr_SetString(PyExc_ValueError, "n must be at least one");
148 : 0 : return NULL;
149 : : }
150 : 0 : it = PyObject_GetIter(iterable);
151 [ # # ]: 0 : if (it == NULL) {
152 : 0 : return NULL;
153 : : }
154 : :
155 : : /* create batchedobject structure */
156 : 0 : bo = (batchedobject *)type->tp_alloc(type, 0);
157 [ # # ]: 0 : if (bo == NULL) {
158 : 0 : Py_DECREF(it);
159 : 0 : return NULL;
160 : : }
161 : 0 : bo->batch_size = n;
162 : 0 : bo->it = it;
163 : 0 : return (PyObject *)bo;
164 : : }
165 : :
166 : : static void
167 : 0 : batched_dealloc(batchedobject *bo)
168 : : {
169 : 0 : PyTypeObject *tp = Py_TYPE(bo);
170 : 0 : PyObject_GC_UnTrack(bo);
171 : 0 : Py_XDECREF(bo->it);
172 : 0 : tp->tp_free(bo);
173 : 0 : Py_DECREF(tp);
174 : 0 : }
175 : :
176 : : static int
177 : 0 : batched_traverse(batchedobject *bo, visitproc visit, void *arg)
178 : : {
179 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(bo));
180 [ # # # # ]: 0 : Py_VISIT(bo->it);
181 : 0 : return 0;
182 : : }
183 : :
184 : : static PyObject *
185 : 0 : batched_next(batchedobject *bo)
186 : : {
187 : : Py_ssize_t i;
188 : 0 : Py_ssize_t n = bo->batch_size;
189 : 0 : PyObject *it = bo->it;
190 : : PyObject *item;
191 : : PyObject *result;
192 : :
193 [ # # ]: 0 : if (it == NULL) {
194 : 0 : return NULL;
195 : : }
196 : 0 : result = PyTuple_New(n);
197 [ # # ]: 0 : if (result == NULL) {
198 : 0 : return NULL;
199 : : }
200 : 0 : iternextfunc iternext = *Py_TYPE(it)->tp_iternext;
201 : 0 : PyObject **items = _PyTuple_ITEMS(result);
202 [ # # ]: 0 : for (i=0 ; i < n ; i++) {
203 : 0 : item = iternext(it);
204 [ # # ]: 0 : if (item == NULL) {
205 : 0 : goto null_item;
206 : : }
207 : 0 : items[i] = item;
208 : : }
209 : 0 : return result;
210 : :
211 : 0 : null_item:
212 [ # # ]: 0 : if (PyErr_Occurred()) {
213 [ # # ]: 0 : if (!PyErr_ExceptionMatches(PyExc_StopIteration)) {
214 : : /* Input raised an exception other than StopIteration */
215 [ # # ]: 0 : Py_CLEAR(bo->it);
216 : 0 : Py_DECREF(result);
217 : 0 : return NULL;
218 : : }
219 : 0 : PyErr_Clear();
220 : : }
221 [ # # ]: 0 : if (i == 0) {
222 [ # # ]: 0 : Py_CLEAR(bo->it);
223 : 0 : Py_DECREF(result);
224 : 0 : return NULL;
225 : : }
226 : 0 : _PyTuple_Resize(&result, i);
227 : 0 : return result;
228 : : }
229 : :
230 : : static PyType_Slot batched_slots[] = {
231 : : {Py_tp_dealloc, batched_dealloc},
232 : : {Py_tp_getattro, PyObject_GenericGetAttr},
233 : : {Py_tp_doc, (void *)batched_new__doc__},
234 : : {Py_tp_traverse, batched_traverse},
235 : : {Py_tp_iter, PyObject_SelfIter},
236 : : {Py_tp_iternext, batched_next},
237 : : {Py_tp_alloc, PyType_GenericAlloc},
238 : : {Py_tp_new, batched_new},
239 : : {Py_tp_free, PyObject_GC_Del},
240 : : {0, NULL},
241 : : };
242 : :
243 : : static PyType_Spec batched_spec = {
244 : : .name = "itertools.batched",
245 : : .basicsize = sizeof(batchedobject),
246 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
247 : : Py_TPFLAGS_IMMUTABLETYPE),
248 : : .slots = batched_slots,
249 : : };
250 : :
251 : :
252 : : /* pairwise object ***********************************************************/
253 : :
254 : : typedef struct {
255 : : PyObject_HEAD
256 : : PyObject *it;
257 : : PyObject *old;
258 : : } pairwiseobject;
259 : :
260 : : /*[clinic input]
261 : : @classmethod
262 : : itertools.pairwise.__new__ as pairwise_new
263 : : iterable: object
264 : : /
265 : : Return an iterator of overlapping pairs taken from the input iterator.
266 : :
267 : : s -> (s0,s1), (s1,s2), (s2, s3), ...
268 : :
269 : : [clinic start generated code]*/
270 : :
271 : : static PyObject *
272 : 0 : pairwise_new_impl(PyTypeObject *type, PyObject *iterable)
273 : : /*[clinic end generated code: output=9f0267062d384456 input=6e7c3cddb431a8d6]*/
274 : : {
275 : : PyObject *it;
276 : : pairwiseobject *po;
277 : :
278 : 0 : it = PyObject_GetIter(iterable);
279 [ # # ]: 0 : if (it == NULL) {
280 : 0 : return NULL;
281 : : }
282 : 0 : po = (pairwiseobject *)type->tp_alloc(type, 0);
283 [ # # ]: 0 : if (po == NULL) {
284 : 0 : Py_DECREF(it);
285 : 0 : return NULL;
286 : : }
287 : 0 : po->it = it;
288 : 0 : po->old = NULL;
289 : 0 : return (PyObject *)po;
290 : : }
291 : :
292 : : static void
293 : 0 : pairwise_dealloc(pairwiseobject *po)
294 : : {
295 : 0 : PyTypeObject *tp = Py_TYPE(po);
296 : 0 : PyObject_GC_UnTrack(po);
297 : 0 : Py_XDECREF(po->it);
298 : 0 : Py_XDECREF(po->old);
299 : 0 : tp->tp_free(po);
300 : 0 : Py_DECREF(tp);
301 : 0 : }
302 : :
303 : : static int
304 : 0 : pairwise_traverse(pairwiseobject *po, visitproc visit, void *arg)
305 : : {
306 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(po));
307 [ # # # # ]: 0 : Py_VISIT(po->it);
308 [ # # # # ]: 0 : Py_VISIT(po->old);
309 : 0 : return 0;
310 : : }
311 : :
312 : : static PyObject *
313 : 0 : pairwise_next(pairwiseobject *po)
314 : : {
315 : 0 : PyObject *it = po->it;
316 : 0 : PyObject *old = po->old;
317 : : PyObject *new, *result;
318 : :
319 [ # # ]: 0 : if (it == NULL) {
320 : 0 : return NULL;
321 : : }
322 [ # # ]: 0 : if (old == NULL) {
323 : 0 : po->old = old = (*Py_TYPE(it)->tp_iternext)(it);
324 [ # # ]: 0 : if (old == NULL) {
325 [ # # ]: 0 : Py_CLEAR(po->it);
326 : 0 : return NULL;
327 : : }
328 : : }
329 : 0 : new = (*Py_TYPE(it)->tp_iternext)(it);
330 [ # # ]: 0 : if (new == NULL) {
331 [ # # ]: 0 : Py_CLEAR(po->it);
332 [ # # ]: 0 : Py_CLEAR(po->old);
333 : 0 : return NULL;
334 : : }
335 : : /* Future optimization: Reuse the result tuple as we do in enumerate() */
336 : 0 : result = PyTuple_Pack(2, old, new);
337 : 0 : Py_SETREF(po->old, new);
338 : 0 : return result;
339 : : }
340 : :
341 : : static PyType_Slot pairwise_slots[] = {
342 : : {Py_tp_dealloc, pairwise_dealloc},
343 : : {Py_tp_getattro, PyObject_GenericGetAttr},
344 : : {Py_tp_doc, (void *)pairwise_new__doc__},
345 : : {Py_tp_traverse, pairwise_traverse},
346 : : {Py_tp_iter, PyObject_SelfIter},
347 : : {Py_tp_iternext, pairwise_next},
348 : : {Py_tp_alloc, PyType_GenericAlloc},
349 : : {Py_tp_new, pairwise_new},
350 : : {Py_tp_free, PyObject_GC_Del},
351 : : {0, NULL},
352 : : };
353 : :
354 : : static PyType_Spec pairwise_spec = {
355 : : .name = "itertools.pairwise",
356 : : .basicsize = sizeof(pairwiseobject),
357 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
358 : : Py_TPFLAGS_IMMUTABLETYPE),
359 : : .slots = pairwise_slots,
360 : : };
361 : :
362 : :
363 : : /* groupby object ************************************************************/
364 : :
365 : : typedef struct {
366 : : PyObject_HEAD
367 : : PyObject *it;
368 : : PyObject *keyfunc;
369 : : PyObject *tgtkey;
370 : : PyObject *currkey;
371 : : PyObject *currvalue;
372 : : const void *currgrouper; /* borrowed reference */
373 : : itertools_state *state;
374 : : } groupbyobject;
375 : :
376 : : static PyObject *_grouper_create(groupbyobject *, PyObject *);
377 : :
378 : : /*[clinic input]
379 : : @classmethod
380 : : itertools.groupby.__new__
381 : :
382 : : iterable as it: object
383 : : Elements to divide into groups according to the key function.
384 : : key as keyfunc: object = None
385 : : A function for computing the group category for each element.
386 : : If the key function is not specified or is None, the element itself
387 : : is used for grouping.
388 : :
389 : : make an iterator that returns consecutive keys and groups from the iterable
390 : : [clinic start generated code]*/
391 : :
392 : : static PyObject *
393 : 0 : itertools_groupby_impl(PyTypeObject *type, PyObject *it, PyObject *keyfunc)
394 : : /*[clinic end generated code: output=cbb1ae3a90fd4141 input=6b3d123e87ff65a1]*/
395 : : {
396 : : groupbyobject *gbo;
397 : :
398 : 0 : gbo = (groupbyobject *)type->tp_alloc(type, 0);
399 [ # # ]: 0 : if (gbo == NULL)
400 : 0 : return NULL;
401 : 0 : gbo->tgtkey = NULL;
402 : 0 : gbo->currkey = NULL;
403 : 0 : gbo->currvalue = NULL;
404 : 0 : gbo->keyfunc = Py_NewRef(keyfunc);
405 : 0 : gbo->it = PyObject_GetIter(it);
406 [ # # ]: 0 : if (gbo->it == NULL) {
407 : 0 : Py_DECREF(gbo);
408 : 0 : return NULL;
409 : : }
410 : 0 : gbo->state = find_state_by_type(type);
411 : 0 : return (PyObject *)gbo;
412 : : }
413 : :
414 : : static void
415 : 0 : groupby_dealloc(groupbyobject *gbo)
416 : : {
417 : 0 : PyTypeObject *tp = Py_TYPE(gbo);
418 : 0 : PyObject_GC_UnTrack(gbo);
419 : 0 : Py_XDECREF(gbo->it);
420 : 0 : Py_XDECREF(gbo->keyfunc);
421 : 0 : Py_XDECREF(gbo->tgtkey);
422 : 0 : Py_XDECREF(gbo->currkey);
423 : 0 : Py_XDECREF(gbo->currvalue);
424 : 0 : tp->tp_free(gbo);
425 : 0 : Py_DECREF(tp);
426 : 0 : }
427 : :
428 : : static int
429 : 0 : groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
430 : : {
431 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(gbo));
432 [ # # # # ]: 0 : Py_VISIT(gbo->it);
433 [ # # # # ]: 0 : Py_VISIT(gbo->keyfunc);
434 [ # # # # ]: 0 : Py_VISIT(gbo->tgtkey);
435 [ # # # # ]: 0 : Py_VISIT(gbo->currkey);
436 [ # # # # ]: 0 : Py_VISIT(gbo->currvalue);
437 : 0 : return 0;
438 : : }
439 : :
440 : : Py_LOCAL_INLINE(int)
441 : 0 : groupby_step(groupbyobject *gbo)
442 : : {
443 : : PyObject *newvalue, *newkey, *oldvalue;
444 : :
445 : 0 : newvalue = PyIter_Next(gbo->it);
446 [ # # ]: 0 : if (newvalue == NULL)
447 : 0 : return -1;
448 : :
449 [ # # ]: 0 : if (gbo->keyfunc == Py_None) {
450 : 0 : newkey = Py_NewRef(newvalue);
451 : : } else {
452 : 0 : newkey = PyObject_CallOneArg(gbo->keyfunc, newvalue);
453 [ # # ]: 0 : if (newkey == NULL) {
454 : 0 : Py_DECREF(newvalue);
455 : 0 : return -1;
456 : : }
457 : : }
458 : :
459 : 0 : oldvalue = gbo->currvalue;
460 : 0 : gbo->currvalue = newvalue;
461 : 0 : Py_XSETREF(gbo->currkey, newkey);
462 : 0 : Py_XDECREF(oldvalue);
463 : 0 : return 0;
464 : : }
465 : :
466 : : static PyObject *
467 : 0 : groupby_next(groupbyobject *gbo)
468 : : {
469 : : PyObject *r, *grouper;
470 : :
471 : 0 : gbo->currgrouper = NULL;
472 : : /* skip to next iteration group */
473 : : for (;;) {
474 [ # # ]: 0 : if (gbo->currkey == NULL)
475 : : /* pass */;
476 [ # # ]: 0 : else if (gbo->tgtkey == NULL)
477 : 0 : break;
478 : : else {
479 : : int rcmp;
480 : :
481 : 0 : rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ);
482 [ # # ]: 0 : if (rcmp == -1)
483 : 0 : return NULL;
484 [ # # ]: 0 : else if (rcmp == 0)
485 : 0 : break;
486 : : }
487 : :
488 [ # # ]: 0 : if (groupby_step(gbo) < 0)
489 : 0 : return NULL;
490 : : }
491 : 0 : Py_INCREF(gbo->currkey);
492 : 0 : Py_XSETREF(gbo->tgtkey, gbo->currkey);
493 : :
494 : 0 : grouper = _grouper_create(gbo, gbo->tgtkey);
495 [ # # ]: 0 : if (grouper == NULL)
496 : 0 : return NULL;
497 : :
498 : 0 : r = PyTuple_Pack(2, gbo->currkey, grouper);
499 : 0 : Py_DECREF(grouper);
500 : 0 : return r;
501 : : }
502 : :
503 : : static PyObject *
504 : 0 : groupby_reduce(groupbyobject *lz, PyObject *Py_UNUSED(ignored))
505 : : {
506 : : /* reduce as a 'new' call with an optional 'setstate' if groupby
507 : : * has started
508 : : */
509 : : PyObject *value;
510 [ # # # # : 0 : if (lz->tgtkey && lz->currkey && lz->currvalue)
# # ]
511 : 0 : value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
512 : : lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
513 : : else
514 : 0 : value = Py_BuildValue("O(OO)", Py_TYPE(lz),
515 : : lz->it, lz->keyfunc);
516 : :
517 : 0 : return value;
518 : : }
519 : :
520 : : PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
521 : :
522 : : static PyObject *
523 : 0 : groupby_setstate(groupbyobject *lz, PyObject *state)
524 : : {
525 : : PyObject *currkey, *currvalue, *tgtkey;
526 [ # # ]: 0 : if (!PyTuple_Check(state)) {
527 : 0 : PyErr_SetString(PyExc_TypeError, "state is not a tuple");
528 : 0 : return NULL;
529 : : }
530 [ # # ]: 0 : if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey)) {
531 : 0 : return NULL;
532 : : }
533 : 0 : Py_INCREF(currkey);
534 : 0 : Py_XSETREF(lz->currkey, currkey);
535 : 0 : Py_INCREF(currvalue);
536 : 0 : Py_XSETREF(lz->currvalue, currvalue);
537 : 0 : Py_INCREF(tgtkey);
538 : 0 : Py_XSETREF(lz->tgtkey, tgtkey);
539 : 0 : Py_RETURN_NONE;
540 : : }
541 : :
542 : : PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
543 : :
544 : : static PyMethodDef groupby_methods[] = {
545 : : {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
546 : : reduce_doc},
547 : : {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
548 : : setstate_doc},
549 : : {NULL, NULL} /* sentinel */
550 : : };
551 : :
552 : : static PyType_Slot groupby_slots[] = {
553 : : {Py_tp_dealloc, groupby_dealloc},
554 : : {Py_tp_getattro, PyObject_GenericGetAttr},
555 : : {Py_tp_doc, (void *)itertools_groupby__doc__},
556 : : {Py_tp_traverse, groupby_traverse},
557 : : {Py_tp_iter, PyObject_SelfIter},
558 : : {Py_tp_iternext, groupby_next},
559 : : {Py_tp_methods, groupby_methods},
560 : : {Py_tp_new, itertools_groupby},
561 : : {Py_tp_free, PyObject_GC_Del},
562 : : {0, NULL},
563 : : };
564 : :
565 : : static PyType_Spec groupby_spec = {
566 : : .name = "itertools.groupby",
567 : : .basicsize= sizeof(groupbyobject),
568 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
569 : : Py_TPFLAGS_IMMUTABLETYPE),
570 : : .slots = groupby_slots,
571 : : };
572 : :
573 : : /* _grouper object (internal) ************************************************/
574 : :
575 : : typedef struct {
576 : : PyObject_HEAD
577 : : PyObject *parent;
578 : : PyObject *tgtkey;
579 : : } _grouperobject;
580 : :
581 : : /*[clinic input]
582 : : @classmethod
583 : : itertools._grouper.__new__
584 : :
585 : : parent: object(subclass_of='clinic_state_by_cls()->groupby_type')
586 : : tgtkey: object
587 : : /
588 : : [clinic start generated code]*/
589 : :
590 : : static PyObject *
591 : 0 : itertools__grouper_impl(PyTypeObject *type, PyObject *parent,
592 : : PyObject *tgtkey)
593 : : /*[clinic end generated code: output=462efb1cdebb5914 input=afe05eb477118f12]*/
594 : : {
595 : 0 : return _grouper_create((groupbyobject*) parent, tgtkey);
596 : : }
597 : :
598 : : static PyObject *
599 : 0 : _grouper_create(groupbyobject *parent, PyObject *tgtkey)
600 : : {
601 : 0 : itertools_state *state = parent->state;
602 : 0 : _grouperobject *igo = PyObject_GC_New(_grouperobject, state->_grouper_type);
603 [ # # ]: 0 : if (igo == NULL)
604 : 0 : return NULL;
605 : 0 : igo->parent = Py_NewRef(parent);
606 : 0 : igo->tgtkey = Py_NewRef(tgtkey);
607 : 0 : parent->currgrouper = igo; /* borrowed reference */
608 : :
609 : 0 : PyObject_GC_Track(igo);
610 : 0 : return (PyObject *)igo;
611 : : }
612 : :
613 : : static void
614 : 0 : _grouper_dealloc(_grouperobject *igo)
615 : : {
616 : 0 : PyTypeObject *tp = Py_TYPE(igo);
617 : 0 : PyObject_GC_UnTrack(igo);
618 : 0 : Py_DECREF(igo->parent);
619 : 0 : Py_DECREF(igo->tgtkey);
620 : 0 : PyObject_GC_Del(igo);
621 : 0 : Py_DECREF(tp);
622 : 0 : }
623 : :
624 : : static int
625 : 0 : _grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
626 : : {
627 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(igo));
628 [ # # # # ]: 0 : Py_VISIT(igo->parent);
629 [ # # # # ]: 0 : Py_VISIT(igo->tgtkey);
630 : 0 : return 0;
631 : : }
632 : :
633 : : static PyObject *
634 : 0 : _grouper_next(_grouperobject *igo)
635 : : {
636 : 0 : groupbyobject *gbo = (groupbyobject *)igo->parent;
637 : : PyObject *r;
638 : : int rcmp;
639 : :
640 [ # # ]: 0 : if (gbo->currgrouper != igo)
641 : 0 : return NULL;
642 [ # # ]: 0 : if (gbo->currvalue == NULL) {
643 [ # # ]: 0 : if (groupby_step(gbo) < 0)
644 : 0 : return NULL;
645 : : }
646 : :
647 : : assert(gbo->currkey != NULL);
648 : 0 : rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
649 [ # # ]: 0 : if (rcmp <= 0)
650 : : /* got any error or current group is end */
651 : 0 : return NULL;
652 : :
653 : 0 : r = gbo->currvalue;
654 : 0 : gbo->currvalue = NULL;
655 [ # # ]: 0 : Py_CLEAR(gbo->currkey);
656 : :
657 : 0 : return r;
658 : : }
659 : :
660 : : static PyObject *
661 : 0 : _grouper_reduce(_grouperobject *lz, PyObject *Py_UNUSED(ignored))
662 : : {
663 [ # # ]: 0 : if (((groupbyobject *)lz->parent)->currgrouper != lz) {
664 : 0 : return Py_BuildValue("N(())", _PyEval_GetBuiltin(&_Py_ID(iter)));
665 : : }
666 : 0 : return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey);
667 : : }
668 : :
669 : : static PyMethodDef _grouper_methods[] = {
670 : : {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
671 : : reduce_doc},
672 : : {NULL, NULL} /* sentinel */
673 : : };
674 : :
675 : : static PyType_Slot _grouper_slots[] = {
676 : : {Py_tp_dealloc, _grouper_dealloc},
677 : : {Py_tp_getattro, PyObject_GenericGetAttr},
678 : : {Py_tp_traverse, _grouper_traverse},
679 : : {Py_tp_iter, PyObject_SelfIter},
680 : : {Py_tp_iternext, _grouper_next},
681 : : {Py_tp_methods, _grouper_methods},
682 : : {Py_tp_new, itertools__grouper},
683 : : {Py_tp_free, PyObject_GC_Del},
684 : : {0, NULL},
685 : : };
686 : :
687 : : static PyType_Spec _grouper_spec = {
688 : : .name = "itertools._grouper",
689 : : .basicsize = sizeof(_grouperobject),
690 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
691 : : Py_TPFLAGS_IMMUTABLETYPE),
692 : : .slots = _grouper_slots,
693 : : };
694 : :
695 : :
696 : : /* tee object and with supporting function and objects ***********************/
697 : :
698 : : /* The teedataobject pre-allocates space for LINKCELLS number of objects.
699 : : To help the object fit neatly inside cache lines (space for 16 to 32
700 : : pointers), the value should be a multiple of 16 minus space for
701 : : the other structure members including PyHEAD overhead. The larger the
702 : : value, the less memory overhead per object and the less time spent
703 : : allocating/deallocating new links. The smaller the number, the less
704 : : wasted space and the more rapid freeing of older data.
705 : : */
706 : : #define LINKCELLS 57
707 : :
708 : : typedef struct {
709 : : PyObject_HEAD
710 : : PyObject *it;
711 : : int numread; /* 0 <= numread <= LINKCELLS */
712 : : int running;
713 : : PyObject *nextlink;
714 : : PyObject *(values[LINKCELLS]);
715 : : } teedataobject;
716 : :
717 : : typedef struct {
718 : : PyObject_HEAD
719 : : teedataobject *dataobj;
720 : : int index; /* 0 <= index <= LINKCELLS */
721 : : PyObject *weakreflist;
722 : : itertools_state *state;
723 : : } teeobject;
724 : :
725 : : static PyObject *
726 : 0 : teedataobject_newinternal(itertools_state *state, PyObject *it)
727 : : {
728 : : teedataobject *tdo;
729 : :
730 : 0 : tdo = PyObject_GC_New(teedataobject, state->teedataobject_type);
731 [ # # ]: 0 : if (tdo == NULL)
732 : 0 : return NULL;
733 : :
734 : 0 : tdo->running = 0;
735 : 0 : tdo->numread = 0;
736 : 0 : tdo->nextlink = NULL;
737 : 0 : tdo->it = Py_NewRef(it);
738 : 0 : PyObject_GC_Track(tdo);
739 : 0 : return (PyObject *)tdo;
740 : : }
741 : :
742 : : static PyObject *
743 : 0 : teedataobject_jumplink(itertools_state *state, teedataobject *tdo)
744 : : {
745 [ # # ]: 0 : if (tdo->nextlink == NULL)
746 : 0 : tdo->nextlink = teedataobject_newinternal(state, tdo->it);
747 : 0 : return Py_XNewRef(tdo->nextlink);
748 : : }
749 : :
750 : : static PyObject *
751 : 0 : teedataobject_getitem(teedataobject *tdo, int i)
752 : : {
753 : : PyObject *value;
754 : :
755 : : assert(i < LINKCELLS);
756 [ # # ]: 0 : if (i < tdo->numread)
757 : 0 : value = tdo->values[i];
758 : : else {
759 : : /* this is the lead iterator, so fetch more data */
760 : : assert(i == tdo->numread);
761 [ # # ]: 0 : if (tdo->running) {
762 : 0 : PyErr_SetString(PyExc_RuntimeError,
763 : : "cannot re-enter the tee iterator");
764 : 0 : return NULL;
765 : : }
766 : 0 : tdo->running = 1;
767 : 0 : value = PyIter_Next(tdo->it);
768 : 0 : tdo->running = 0;
769 [ # # ]: 0 : if (value == NULL)
770 : 0 : return NULL;
771 : 0 : tdo->numread++;
772 : 0 : tdo->values[i] = value;
773 : : }
774 : 0 : return Py_NewRef(value);
775 : : }
776 : :
777 : : static int
778 : 0 : teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
779 : : {
780 : : int i;
781 : :
782 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(tdo));
783 [ # # # # ]: 0 : Py_VISIT(tdo->it);
784 [ # # ]: 0 : for (i = 0; i < tdo->numread; i++)
785 [ # # # # ]: 0 : Py_VISIT(tdo->values[i]);
786 [ # # # # ]: 0 : Py_VISIT(tdo->nextlink);
787 : 0 : return 0;
788 : : }
789 : :
790 : : static void
791 : 0 : teedataobject_safe_decref(PyObject *obj, PyTypeObject *tdo_type)
792 : : {
793 [ # # # # : 0 : while (obj && Py_IS_TYPE(obj, tdo_type) &&
# # ]
794 : 0 : Py_REFCNT(obj) == 1) {
795 : 0 : PyObject *nextlink = ((teedataobject *)obj)->nextlink;
796 : 0 : ((teedataobject *)obj)->nextlink = NULL;
797 : 0 : Py_SETREF(obj, nextlink);
798 : : }
799 : 0 : Py_XDECREF(obj);
800 : 0 : }
801 : :
802 : : static int
803 : 0 : teedataobject_clear(teedataobject *tdo)
804 : : {
805 : : int i;
806 : : PyObject *tmp;
807 : :
808 [ # # ]: 0 : Py_CLEAR(tdo->it);
809 [ # # ]: 0 : for (i=0 ; i<tdo->numread ; i++)
810 [ # # ]: 0 : Py_CLEAR(tdo->values[i]);
811 : 0 : tmp = tdo->nextlink;
812 : 0 : tdo->nextlink = NULL;
813 : 0 : itertools_state *state = get_module_state_by_cls(Py_TYPE(tdo));
814 : 0 : teedataobject_safe_decref(tmp, state->teedataobject_type);
815 : 0 : return 0;
816 : : }
817 : :
818 : : static void
819 : 0 : teedataobject_dealloc(teedataobject *tdo)
820 : : {
821 : 0 : PyTypeObject *tp = Py_TYPE(tdo);
822 : 0 : PyObject_GC_UnTrack(tdo);
823 : 0 : teedataobject_clear(tdo);
824 : 0 : PyObject_GC_Del(tdo);
825 : 0 : Py_DECREF(tp);
826 : 0 : }
827 : :
828 : : static PyObject *
829 : 0 : teedataobject_reduce(teedataobject *tdo, PyObject *Py_UNUSED(ignored))
830 : : {
831 : : int i;
832 : : /* create a temporary list of already iterated values */
833 : 0 : PyObject *values = PyList_New(tdo->numread);
834 : :
835 [ # # ]: 0 : if (!values)
836 : 0 : return NULL;
837 [ # # ]: 0 : for (i=0 ; i<tdo->numread ; i++) {
838 : 0 : Py_INCREF(tdo->values[i]);
839 : 0 : PyList_SET_ITEM(values, i, tdo->values[i]);
840 : : }
841 : 0 : return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
842 : : values,
843 [ # # ]: 0 : tdo->nextlink ? tdo->nextlink : Py_None);
844 : : }
845 : :
846 : : /*[clinic input]
847 : : @classmethod
848 : : itertools.teedataobject.__new__
849 : : iterable as it: object
850 : : values: object(subclass_of='&PyList_Type')
851 : : next: object
852 : : /
853 : : Data container common to multiple tee objects.
854 : : [clinic start generated code]*/
855 : :
856 : : static PyObject *
857 : 0 : itertools_teedataobject_impl(PyTypeObject *type, PyObject *it,
858 : : PyObject *values, PyObject *next)
859 : : /*[clinic end generated code: output=3343ceb07e08df5e input=be60f2fabd2b72ba]*/
860 : : {
861 : : teedataobject *tdo;
862 : : Py_ssize_t i, len;
863 : :
864 : 0 : itertools_state *state = get_module_state_by_cls(type);
865 : : assert(type == state->teedataobject_type);
866 : :
867 : 0 : tdo = (teedataobject *)teedataobject_newinternal(state, it);
868 [ # # ]: 0 : if (!tdo)
869 : 0 : return NULL;
870 : :
871 : 0 : len = PyList_GET_SIZE(values);
872 [ # # ]: 0 : if (len > LINKCELLS)
873 : 0 : goto err;
874 [ # # ]: 0 : for (i=0; i<len; i++) {
875 : 0 : tdo->values[i] = PyList_GET_ITEM(values, i);
876 : 0 : Py_INCREF(tdo->values[i]);
877 : : }
878 : : /* len <= LINKCELLS < INT_MAX */
879 : 0 : tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
880 : :
881 [ # # ]: 0 : if (len == LINKCELLS) {
882 [ # # ]: 0 : if (next != Py_None) {
883 [ # # ]: 0 : if (!Py_IS_TYPE(next, state->teedataobject_type))
884 : 0 : goto err;
885 : : assert(tdo->nextlink == NULL);
886 : 0 : tdo->nextlink = Py_NewRef(next);
887 : : }
888 : : } else {
889 [ # # ]: 0 : if (next != Py_None)
890 : 0 : goto err; /* shouldn't have a next if we are not full */
891 : : }
892 : 0 : return (PyObject*)tdo;
893 : :
894 : 0 : err:
895 : 0 : Py_XDECREF(tdo);
896 : 0 : PyErr_SetString(PyExc_ValueError, "Invalid arguments");
897 : 0 : return NULL;
898 : : }
899 : :
900 : : static PyMethodDef teedataobject_methods[] = {
901 : : {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
902 : : reduce_doc},
903 : : {NULL, NULL} /* sentinel */
904 : : };
905 : :
906 : : static PyType_Slot teedataobject_slots[] = {
907 : : {Py_tp_dealloc, teedataobject_dealloc},
908 : : {Py_tp_getattro, PyObject_GenericGetAttr},
909 : : {Py_tp_doc, (void *)itertools_teedataobject__doc__},
910 : : {Py_tp_traverse, teedataobject_traverse},
911 : : {Py_tp_clear, teedataobject_clear},
912 : : {Py_tp_methods, teedataobject_methods},
913 : : {Py_tp_new, itertools_teedataobject},
914 : : {Py_tp_free, PyObject_GC_Del},
915 : : {0, NULL},
916 : : };
917 : :
918 : : static PyType_Spec teedataobject_spec = {
919 : : .name = "itertools._tee_dataobject",
920 : : .basicsize = sizeof(teedataobject),
921 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
922 : : Py_TPFLAGS_IMMUTABLETYPE),
923 : : .slots = teedataobject_slots,
924 : : };
925 : :
926 : :
927 : : static PyObject *
928 : 0 : tee_next(teeobject *to)
929 : : {
930 : : PyObject *value, *link;
931 : :
932 [ # # ]: 0 : if (to->index >= LINKCELLS) {
933 : 0 : link = teedataobject_jumplink(to->state, to->dataobj);
934 [ # # ]: 0 : if (link == NULL)
935 : 0 : return NULL;
936 : 0 : Py_SETREF(to->dataobj, (teedataobject *)link);
937 : 0 : to->index = 0;
938 : : }
939 : 0 : value = teedataobject_getitem(to->dataobj, to->index);
940 [ # # ]: 0 : if (value == NULL)
941 : 0 : return NULL;
942 : 0 : to->index++;
943 : 0 : return value;
944 : : }
945 : :
946 : : static int
947 : 0 : tee_traverse(teeobject *to, visitproc visit, void *arg)
948 : : {
949 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(to));
950 [ # # # # ]: 0 : Py_VISIT((PyObject *)to->dataobj);
951 : 0 : return 0;
952 : : }
953 : :
954 : : static PyObject *
955 : 0 : tee_copy(teeobject *to, PyObject *Py_UNUSED(ignored))
956 : : {
957 : : teeobject *newto;
958 : :
959 : 0 : newto = PyObject_GC_New(teeobject, Py_TYPE(to));
960 [ # # ]: 0 : if (newto == NULL)
961 : 0 : return NULL;
962 : 0 : newto->dataobj = (teedataobject*)Py_NewRef(to->dataobj);
963 : 0 : newto->index = to->index;
964 : 0 : newto->weakreflist = NULL;
965 : 0 : newto->state = to->state;
966 : 0 : PyObject_GC_Track(newto);
967 : 0 : return (PyObject *)newto;
968 : : }
969 : :
970 : : PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
971 : :
972 : : static PyObject *
973 : 0 : tee_fromiterable(itertools_state *state, PyObject *iterable)
974 : : {
975 : : teeobject *to;
976 : : PyObject *it;
977 : :
978 : 0 : it = PyObject_GetIter(iterable);
979 [ # # ]: 0 : if (it == NULL)
980 : 0 : return NULL;
981 [ # # ]: 0 : if (PyObject_TypeCheck(it, state->tee_type)) {
982 : 0 : to = (teeobject *)tee_copy((teeobject *)it, NULL);
983 : 0 : goto done;
984 : : }
985 : :
986 : 0 : PyObject *dataobj = teedataobject_newinternal(state, it);
987 [ # # ]: 0 : if (!dataobj) {
988 : 0 : to = NULL;
989 : 0 : goto done;
990 : : }
991 : 0 : to = PyObject_GC_New(teeobject, state->tee_type);
992 [ # # ]: 0 : if (to == NULL) {
993 : 0 : Py_DECREF(dataobj);
994 : 0 : goto done;
995 : : }
996 : 0 : to->dataobj = (teedataobject *)dataobj;
997 : 0 : to->index = 0;
998 : 0 : to->weakreflist = NULL;
999 : 0 : to->state = state;
1000 : 0 : PyObject_GC_Track(to);
1001 : 0 : done:
1002 : 0 : Py_DECREF(it);
1003 : 0 : return (PyObject *)to;
1004 : : }
1005 : :
1006 : : /*[clinic input]
1007 : : @classmethod
1008 : : itertools._tee.__new__
1009 : : iterable: object
1010 : : /
1011 : : Iterator wrapped to make it copyable.
1012 : : [clinic start generated code]*/
1013 : :
1014 : : static PyObject *
1015 : 0 : itertools__tee_impl(PyTypeObject *type, PyObject *iterable)
1016 : : /*[clinic end generated code: output=b02d3fd26c810c3f input=adc0779d2afe37a2]*/
1017 : : {
1018 : 0 : itertools_state *state = get_module_state_by_cls(type);
1019 : 0 : return tee_fromiterable(state, iterable);
1020 : : }
1021 : :
1022 : : static int
1023 : 0 : tee_clear(teeobject *to)
1024 : : {
1025 [ # # ]: 0 : if (to->weakreflist != NULL)
1026 : 0 : PyObject_ClearWeakRefs((PyObject *) to);
1027 [ # # ]: 0 : Py_CLEAR(to->dataobj);
1028 : 0 : return 0;
1029 : : }
1030 : :
1031 : : static void
1032 : 0 : tee_dealloc(teeobject *to)
1033 : : {
1034 : 0 : PyTypeObject *tp = Py_TYPE(to);
1035 : 0 : PyObject_GC_UnTrack(to);
1036 : 0 : tee_clear(to);
1037 : 0 : PyObject_GC_Del(to);
1038 : 0 : Py_DECREF(tp);
1039 : 0 : }
1040 : :
1041 : : static PyObject *
1042 : 0 : tee_reduce(teeobject *to, PyObject *Py_UNUSED(ignored))
1043 : : {
1044 : 0 : return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
1045 : : }
1046 : :
1047 : : static PyObject *
1048 : 0 : tee_setstate(teeobject *to, PyObject *state)
1049 : : {
1050 : : teedataobject *tdo;
1051 : : int index;
1052 [ # # ]: 0 : if (!PyTuple_Check(state)) {
1053 : 0 : PyErr_SetString(PyExc_TypeError, "state is not a tuple");
1054 : 0 : return NULL;
1055 : : }
1056 : 0 : PyTypeObject *tdo_type = to->state->teedataobject_type;
1057 [ # # ]: 0 : if (!PyArg_ParseTuple(state, "O!i", tdo_type, &tdo, &index)) {
1058 : 0 : return NULL;
1059 : : }
1060 [ # # # # ]: 0 : if (index < 0 || index > LINKCELLS) {
1061 : 0 : PyErr_SetString(PyExc_ValueError, "Index out of range");
1062 : 0 : return NULL;
1063 : : }
1064 : 0 : Py_INCREF(tdo);
1065 : 0 : Py_XSETREF(to->dataobj, tdo);
1066 : 0 : to->index = index;
1067 : 0 : Py_RETURN_NONE;
1068 : : }
1069 : :
1070 : : static PyMethodDef tee_methods[] = {
1071 : : {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
1072 : : {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
1073 : : {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
1074 : : {NULL, NULL} /* sentinel */
1075 : : };
1076 : :
1077 : : static PyMemberDef tee_members[] = {
1078 : : {"__weaklistoffset__", T_PYSSIZET, offsetof(teeobject, weakreflist), READONLY},
1079 : : {NULL},
1080 : : };
1081 : :
1082 : : static PyType_Slot tee_slots[] = {
1083 : : {Py_tp_dealloc, tee_dealloc},
1084 : : {Py_tp_doc, (void *)itertools__tee__doc__},
1085 : : {Py_tp_traverse, tee_traverse},
1086 : : {Py_tp_clear, tee_clear},
1087 : : {Py_tp_iter, PyObject_SelfIter},
1088 : : {Py_tp_iternext, tee_next},
1089 : : {Py_tp_methods, tee_methods},
1090 : : {Py_tp_members, tee_members},
1091 : : {Py_tp_new, itertools__tee},
1092 : : {Py_tp_free, PyObject_GC_Del},
1093 : : {0, NULL},
1094 : : };
1095 : :
1096 : : static PyType_Spec tee_spec = {
1097 : : .name = "itertools._tee",
1098 : : .basicsize = sizeof(teeobject),
1099 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1100 : : Py_TPFLAGS_IMMUTABLETYPE),
1101 : : .slots = tee_slots,
1102 : : };
1103 : :
1104 : : /*[clinic input]
1105 : : itertools.tee
1106 : : iterable: object
1107 : : n: Py_ssize_t = 2
1108 : : /
1109 : : Returns a tuple of n independent iterators.
1110 : : [clinic start generated code]*/
1111 : :
1112 : : static PyObject *
1113 : 0 : itertools_tee_impl(PyObject *module, PyObject *iterable, Py_ssize_t n)
1114 : : /*[clinic end generated code: output=1c64519cd859c2f0 input=c99a1472c425d66d]*/
1115 : : {
1116 : : Py_ssize_t i;
1117 : : PyObject *it, *copyable, *copyfunc, *result;
1118 : :
1119 [ # # ]: 0 : if (n < 0) {
1120 : 0 : PyErr_SetString(PyExc_ValueError, "n must be >= 0");
1121 : 0 : return NULL;
1122 : : }
1123 : 0 : result = PyTuple_New(n);
1124 [ # # ]: 0 : if (result == NULL)
1125 : 0 : return NULL;
1126 [ # # ]: 0 : if (n == 0)
1127 : 0 : return result;
1128 : 0 : it = PyObject_GetIter(iterable);
1129 [ # # ]: 0 : if (it == NULL) {
1130 : 0 : Py_DECREF(result);
1131 : 0 : return NULL;
1132 : : }
1133 : :
1134 [ # # ]: 0 : if (_PyObject_LookupAttr(it, &_Py_ID(__copy__), ©func) < 0) {
1135 : 0 : Py_DECREF(it);
1136 : 0 : Py_DECREF(result);
1137 : 0 : return NULL;
1138 : : }
1139 [ # # ]: 0 : if (copyfunc != NULL) {
1140 : 0 : copyable = it;
1141 : : }
1142 : : else {
1143 : 0 : itertools_state *state = get_module_state(module);
1144 : 0 : copyable = tee_fromiterable(state, it);
1145 : 0 : Py_DECREF(it);
1146 [ # # ]: 0 : if (copyable == NULL) {
1147 : 0 : Py_DECREF(result);
1148 : 0 : return NULL;
1149 : : }
1150 : 0 : copyfunc = PyObject_GetAttr(copyable, &_Py_ID(__copy__));
1151 [ # # ]: 0 : if (copyfunc == NULL) {
1152 : 0 : Py_DECREF(copyable);
1153 : 0 : Py_DECREF(result);
1154 : 0 : return NULL;
1155 : : }
1156 : : }
1157 : :
1158 : 0 : PyTuple_SET_ITEM(result, 0, copyable);
1159 [ # # ]: 0 : for (i = 1; i < n; i++) {
1160 : 0 : copyable = _PyObject_CallNoArgs(copyfunc);
1161 [ # # ]: 0 : if (copyable == NULL) {
1162 : 0 : Py_DECREF(copyfunc);
1163 : 0 : Py_DECREF(result);
1164 : 0 : return NULL;
1165 : : }
1166 : 0 : PyTuple_SET_ITEM(result, i, copyable);
1167 : : }
1168 : 0 : Py_DECREF(copyfunc);
1169 : 0 : return result;
1170 : : }
1171 : :
1172 : :
1173 : : /* cycle object **************************************************************/
1174 : :
1175 : : typedef struct {
1176 : : PyObject_HEAD
1177 : : PyObject *it;
1178 : : PyObject *saved;
1179 : : Py_ssize_t index;
1180 : : int firstpass;
1181 : : } cycleobject;
1182 : :
1183 : : /*[clinic input]
1184 : : @classmethod
1185 : : itertools.cycle.__new__
1186 : : iterable: object
1187 : : /
1188 : : Return elements from the iterable until it is exhausted. Then repeat the sequence indefinitely.
1189 : : [clinic start generated code]*/
1190 : :
1191 : : static PyObject *
1192 : 0 : itertools_cycle_impl(PyTypeObject *type, PyObject *iterable)
1193 : : /*[clinic end generated code: output=f60e5ec17a45b35c input=9d1d84bcf66e908b]*/
1194 : : {
1195 : : PyObject *it;
1196 : : PyObject *saved;
1197 : : cycleobject *lz;
1198 : :
1199 : : /* Get iterator. */
1200 : 0 : it = PyObject_GetIter(iterable);
1201 [ # # ]: 0 : if (it == NULL)
1202 : 0 : return NULL;
1203 : :
1204 : 0 : saved = PyList_New(0);
1205 [ # # ]: 0 : if (saved == NULL) {
1206 : 0 : Py_DECREF(it);
1207 : 0 : return NULL;
1208 : : }
1209 : :
1210 : : /* create cycleobject structure */
1211 : 0 : lz = (cycleobject *)type->tp_alloc(type, 0);
1212 [ # # ]: 0 : if (lz == NULL) {
1213 : 0 : Py_DECREF(it);
1214 : 0 : Py_DECREF(saved);
1215 : 0 : return NULL;
1216 : : }
1217 : 0 : lz->it = it;
1218 : 0 : lz->saved = saved;
1219 : 0 : lz->index = 0;
1220 : 0 : lz->firstpass = 0;
1221 : :
1222 : 0 : return (PyObject *)lz;
1223 : : }
1224 : :
1225 : : static void
1226 : 0 : cycle_dealloc(cycleobject *lz)
1227 : : {
1228 : 0 : PyTypeObject *tp = Py_TYPE(lz);
1229 : 0 : PyObject_GC_UnTrack(lz);
1230 : 0 : Py_XDECREF(lz->it);
1231 : 0 : Py_XDECREF(lz->saved);
1232 : 0 : tp->tp_free(lz);
1233 : 0 : Py_DECREF(tp);
1234 : 0 : }
1235 : :
1236 : : static int
1237 : 0 : cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
1238 : : {
1239 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(lz));
1240 [ # # # # ]: 0 : Py_VISIT(lz->it);
1241 [ # # # # ]: 0 : Py_VISIT(lz->saved);
1242 : 0 : return 0;
1243 : : }
1244 : :
1245 : : static PyObject *
1246 : 0 : cycle_next(cycleobject *lz)
1247 : : {
1248 : : PyObject *item;
1249 : :
1250 [ # # ]: 0 : if (lz->it != NULL) {
1251 : 0 : item = PyIter_Next(lz->it);
1252 [ # # ]: 0 : if (item != NULL) {
1253 [ # # ]: 0 : if (lz->firstpass)
1254 : 0 : return item;
1255 [ # # ]: 0 : if (PyList_Append(lz->saved, item)) {
1256 : 0 : Py_DECREF(item);
1257 : 0 : return NULL;
1258 : : }
1259 : 0 : return item;
1260 : : }
1261 : : /* Note: StopIteration is already cleared by PyIter_Next() */
1262 [ # # ]: 0 : if (PyErr_Occurred())
1263 : 0 : return NULL;
1264 [ # # ]: 0 : Py_CLEAR(lz->it);
1265 : : }
1266 [ # # ]: 0 : if (PyList_GET_SIZE(lz->saved) == 0)
1267 : 0 : return NULL;
1268 : 0 : item = PyList_GET_ITEM(lz->saved, lz->index);
1269 : 0 : lz->index++;
1270 [ # # ]: 0 : if (lz->index >= PyList_GET_SIZE(lz->saved))
1271 : 0 : lz->index = 0;
1272 : 0 : return Py_NewRef(item);
1273 : : }
1274 : :
1275 : : static PyObject *
1276 : 0 : cycle_reduce(cycleobject *lz, PyObject *Py_UNUSED(ignored))
1277 : : {
1278 : : /* Create a new cycle with the iterator tuple, then set the saved state */
1279 [ # # ]: 0 : if (lz->it == NULL) {
1280 : 0 : PyObject *it = PyObject_GetIter(lz->saved);
1281 [ # # ]: 0 : if (it == NULL)
1282 : 0 : return NULL;
1283 [ # # ]: 0 : if (lz->index != 0) {
1284 : 0 : PyObject *res = _PyObject_CallMethod(it, &_Py_ID(__setstate__),
1285 : : "n", lz->index);
1286 [ # # ]: 0 : if (res == NULL) {
1287 : 0 : Py_DECREF(it);
1288 : 0 : return NULL;
1289 : : }
1290 : 0 : Py_DECREF(res);
1291 : : }
1292 : 0 : return Py_BuildValue("O(N)(OO)", Py_TYPE(lz), it, lz->saved, Py_True);
1293 : : }
1294 : 0 : return Py_BuildValue("O(O)(OO)", Py_TYPE(lz), lz->it, lz->saved,
1295 [ # # ]: 0 : lz->firstpass ? Py_True : Py_False);
1296 : : }
1297 : :
1298 : : static PyObject *
1299 : 0 : cycle_setstate(cycleobject *lz, PyObject *state)
1300 : : {
1301 : 0 : PyObject *saved=NULL;
1302 : : int firstpass;
1303 [ # # ]: 0 : if (!PyTuple_Check(state)) {
1304 : 0 : PyErr_SetString(PyExc_TypeError, "state is not a tuple");
1305 : 0 : return NULL;
1306 : : }
1307 : : // The second item can be 1/0 in old pickles and True/False in new pickles
1308 [ # # ]: 0 : if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
1309 : 0 : return NULL;
1310 : : }
1311 : 0 : Py_INCREF(saved);
1312 : 0 : Py_XSETREF(lz->saved, saved);
1313 : 0 : lz->firstpass = firstpass != 0;
1314 : 0 : lz->index = 0;
1315 : 0 : Py_RETURN_NONE;
1316 : : }
1317 : :
1318 : : static PyMethodDef cycle_methods[] = {
1319 : : {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
1320 : : reduce_doc},
1321 : : {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
1322 : : setstate_doc},
1323 : : {NULL, NULL} /* sentinel */
1324 : : };
1325 : :
1326 : : static PyType_Slot cycle_slots[] = {
1327 : : {Py_tp_dealloc, cycle_dealloc},
1328 : : {Py_tp_getattro, PyObject_GenericGetAttr},
1329 : : {Py_tp_doc, (void *)itertools_cycle__doc__},
1330 : : {Py_tp_traverse, cycle_traverse},
1331 : : {Py_tp_iter, PyObject_SelfIter},
1332 : : {Py_tp_iternext, cycle_next},
1333 : : {Py_tp_methods, cycle_methods},
1334 : : {Py_tp_new, itertools_cycle},
1335 : : {Py_tp_free, PyObject_GC_Del},
1336 : : {0, NULL},
1337 : : };
1338 : :
1339 : : static PyType_Spec cycle_spec = {
1340 : : .name = "itertools.cycle",
1341 : : .basicsize = sizeof(cycleobject),
1342 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
1343 : : Py_TPFLAGS_IMMUTABLETYPE),
1344 : : .slots = cycle_slots,
1345 : : };
1346 : :
1347 : :
1348 : : /* dropwhile object **********************************************************/
1349 : :
1350 : : typedef struct {
1351 : : PyObject_HEAD
1352 : : PyObject *func;
1353 : : PyObject *it;
1354 : : long start;
1355 : : } dropwhileobject;
1356 : :
1357 : : /*[clinic input]
1358 : : @classmethod
1359 : : itertools.dropwhile.__new__
1360 : : predicate as func: object
1361 : : iterable as seq: object
1362 : : /
1363 : : Drop items from the iterable while predicate(item) is true.
1364 : :
1365 : : Afterwards, return every element until the iterable is exhausted.
1366 : : [clinic start generated code]*/
1367 : :
1368 : : static PyObject *
1369 : 0 : itertools_dropwhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1370 : : /*[clinic end generated code: output=92f9d0d89af149e4 input=d39737147c9f0a26]*/
1371 : : {
1372 : : PyObject *it;
1373 : : dropwhileobject *lz;
1374 : :
1375 : : /* Get iterator. */
1376 : 0 : it = PyObject_GetIter(seq);
1377 [ # # ]: 0 : if (it == NULL)
1378 : 0 : return NULL;
1379 : :
1380 : : /* create dropwhileobject structure */
1381 : 0 : lz = (dropwhileobject *)type->tp_alloc(type, 0);
1382 [ # # ]: 0 : if (lz == NULL) {
1383 : 0 : Py_DECREF(it);
1384 : 0 : return NULL;
1385 : : }
1386 : 0 : lz->func = Py_NewRef(func);
1387 : 0 : lz->it = it;
1388 : 0 : lz->start = 0;
1389 : :
1390 : 0 : return (PyObject *)lz;
1391 : : }
1392 : :
1393 : : static void
1394 : 0 : dropwhile_dealloc(dropwhileobject *lz)
1395 : : {
1396 : 0 : PyTypeObject *tp = Py_TYPE(lz);
1397 : 0 : PyObject_GC_UnTrack(lz);
1398 : 0 : Py_XDECREF(lz->func);
1399 : 0 : Py_XDECREF(lz->it);
1400 : 0 : tp->tp_free(lz);
1401 : 0 : Py_DECREF(tp);
1402 : 0 : }
1403 : :
1404 : : static int
1405 : 0 : dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
1406 : : {
1407 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(lz));
1408 [ # # # # ]: 0 : Py_VISIT(lz->it);
1409 [ # # # # ]: 0 : Py_VISIT(lz->func);
1410 : 0 : return 0;
1411 : : }
1412 : :
1413 : : static PyObject *
1414 : 0 : dropwhile_next(dropwhileobject *lz)
1415 : : {
1416 : : PyObject *item, *good;
1417 : 0 : PyObject *it = lz->it;
1418 : : long ok;
1419 : : PyObject *(*iternext)(PyObject *);
1420 : :
1421 : 0 : iternext = *Py_TYPE(it)->tp_iternext;
1422 : : for (;;) {
1423 : 0 : item = iternext(it);
1424 [ # # ]: 0 : if (item == NULL)
1425 : 0 : return NULL;
1426 [ # # ]: 0 : if (lz->start == 1)
1427 : 0 : return item;
1428 : :
1429 : 0 : good = PyObject_CallOneArg(lz->func, item);
1430 [ # # ]: 0 : if (good == NULL) {
1431 : 0 : Py_DECREF(item);
1432 : 0 : return NULL;
1433 : : }
1434 : 0 : ok = PyObject_IsTrue(good);
1435 : 0 : Py_DECREF(good);
1436 [ # # ]: 0 : if (ok == 0) {
1437 : 0 : lz->start = 1;
1438 : 0 : return item;
1439 : : }
1440 : 0 : Py_DECREF(item);
1441 [ # # ]: 0 : if (ok < 0)
1442 : 0 : return NULL;
1443 : : }
1444 : : }
1445 : :
1446 : : static PyObject *
1447 : 0 : dropwhile_reduce(dropwhileobject *lz, PyObject *Py_UNUSED(ignored))
1448 : : {
1449 : 0 : return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->start);
1450 : : }
1451 : :
1452 : : static PyObject *
1453 : 0 : dropwhile_setstate(dropwhileobject *lz, PyObject *state)
1454 : : {
1455 : 0 : int start = PyObject_IsTrue(state);
1456 [ # # ]: 0 : if (start < 0)
1457 : 0 : return NULL;
1458 : 0 : lz->start = start;
1459 : 0 : Py_RETURN_NONE;
1460 : : }
1461 : :
1462 : : static PyMethodDef dropwhile_methods[] = {
1463 : : {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
1464 : : reduce_doc},
1465 : : {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
1466 : : setstate_doc},
1467 : : {NULL, NULL} /* sentinel */
1468 : : };
1469 : :
1470 : : static PyType_Slot dropwhile_slots[] = {
1471 : : {Py_tp_dealloc, dropwhile_dealloc},
1472 : : {Py_tp_getattro, PyObject_GenericGetAttr},
1473 : : {Py_tp_doc, (void *)itertools_dropwhile__doc__},
1474 : : {Py_tp_traverse, dropwhile_traverse},
1475 : : {Py_tp_iter, PyObject_SelfIter},
1476 : : {Py_tp_iternext, dropwhile_next},
1477 : : {Py_tp_methods, dropwhile_methods},
1478 : : {Py_tp_new, itertools_dropwhile},
1479 : : {Py_tp_free, PyObject_GC_Del},
1480 : : {0, NULL},
1481 : : };
1482 : :
1483 : : static PyType_Spec dropwhile_spec = {
1484 : : .name = "itertools.dropwhile",
1485 : : .basicsize = sizeof(dropwhileobject),
1486 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
1487 : : Py_TPFLAGS_IMMUTABLETYPE),
1488 : : .slots = dropwhile_slots,
1489 : : };
1490 : :
1491 : :
1492 : : /* takewhile object **********************************************************/
1493 : :
1494 : : typedef struct {
1495 : : PyObject_HEAD
1496 : : PyObject *func;
1497 : : PyObject *it;
1498 : : long stop;
1499 : : } takewhileobject;
1500 : :
1501 : : /*[clinic input]
1502 : : @classmethod
1503 : : itertools.takewhile.__new__
1504 : : predicate as func: object
1505 : : iterable as seq: object
1506 : : /
1507 : : Return successive entries from an iterable as long as the predicate evaluates to true for each entry.
1508 : : [clinic start generated code]*/
1509 : :
1510 : : static PyObject *
1511 : 0 : itertools_takewhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1512 : : /*[clinic end generated code: output=bb179ea7864e2ef6 input=ba5255f7519aa119]*/
1513 : : {
1514 : : PyObject *it;
1515 : : takewhileobject *lz;
1516 : :
1517 : : /* Get iterator. */
1518 : 0 : it = PyObject_GetIter(seq);
1519 [ # # ]: 0 : if (it == NULL)
1520 : 0 : return NULL;
1521 : :
1522 : : /* create takewhileobject structure */
1523 : 0 : lz = (takewhileobject *)type->tp_alloc(type, 0);
1524 [ # # ]: 0 : if (lz == NULL) {
1525 : 0 : Py_DECREF(it);
1526 : 0 : return NULL;
1527 : : }
1528 : 0 : lz->func = Py_NewRef(func);
1529 : 0 : lz->it = it;
1530 : 0 : lz->stop = 0;
1531 : :
1532 : 0 : return (PyObject *)lz;
1533 : : }
1534 : :
1535 : : static void
1536 : 0 : takewhile_dealloc(takewhileobject *lz)
1537 : : {
1538 : 0 : PyTypeObject *tp = Py_TYPE(lz);
1539 : 0 : PyObject_GC_UnTrack(lz);
1540 : 0 : Py_XDECREF(lz->func);
1541 : 0 : Py_XDECREF(lz->it);
1542 : 0 : tp->tp_free(lz);
1543 : 0 : Py_DECREF(tp);
1544 : 0 : }
1545 : :
1546 : : static int
1547 : 0 : takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
1548 : : {
1549 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(lz));
1550 [ # # # # ]: 0 : Py_VISIT(lz->it);
1551 [ # # # # ]: 0 : Py_VISIT(lz->func);
1552 : 0 : return 0;
1553 : : }
1554 : :
1555 : : static PyObject *
1556 : 0 : takewhile_next(takewhileobject *lz)
1557 : : {
1558 : : PyObject *item, *good;
1559 : 0 : PyObject *it = lz->it;
1560 : : long ok;
1561 : :
1562 [ # # ]: 0 : if (lz->stop == 1)
1563 : 0 : return NULL;
1564 : :
1565 : 0 : item = (*Py_TYPE(it)->tp_iternext)(it);
1566 [ # # ]: 0 : if (item == NULL)
1567 : 0 : return NULL;
1568 : :
1569 : 0 : good = PyObject_CallOneArg(lz->func, item);
1570 [ # # ]: 0 : if (good == NULL) {
1571 : 0 : Py_DECREF(item);
1572 : 0 : return NULL;
1573 : : }
1574 : 0 : ok = PyObject_IsTrue(good);
1575 : 0 : Py_DECREF(good);
1576 [ # # ]: 0 : if (ok > 0)
1577 : 0 : return item;
1578 : 0 : Py_DECREF(item);
1579 [ # # ]: 0 : if (ok == 0)
1580 : 0 : lz->stop = 1;
1581 : 0 : return NULL;
1582 : : }
1583 : :
1584 : : static PyObject *
1585 : 0 : takewhile_reduce(takewhileobject *lz, PyObject *Py_UNUSED(ignored))
1586 : : {
1587 : 0 : return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->stop);
1588 : : }
1589 : :
1590 : : static PyObject *
1591 : 0 : takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
1592 : : {
1593 : 0 : int stop = PyObject_IsTrue(state);
1594 : :
1595 [ # # ]: 0 : if (stop < 0)
1596 : 0 : return NULL;
1597 : 0 : lz->stop = stop;
1598 : 0 : Py_RETURN_NONE;
1599 : : }
1600 : :
1601 : : static PyMethodDef takewhile_reduce_methods[] = {
1602 : : {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
1603 : : reduce_doc},
1604 : : {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
1605 : : setstate_doc},
1606 : : {NULL, NULL} /* sentinel */
1607 : : };
1608 : :
1609 : : static PyType_Slot takewhile_slots[] = {
1610 : : {Py_tp_dealloc, takewhile_dealloc},
1611 : : {Py_tp_getattro, PyObject_GenericGetAttr},
1612 : : {Py_tp_doc, (void *)itertools_takewhile__doc__},
1613 : : {Py_tp_traverse, takewhile_traverse},
1614 : : {Py_tp_iter, PyObject_SelfIter},
1615 : : {Py_tp_iternext, takewhile_next},
1616 : : {Py_tp_methods, takewhile_reduce_methods},
1617 : : {Py_tp_new, itertools_takewhile},
1618 : : {Py_tp_free, PyObject_GC_Del},
1619 : : {0, NULL},
1620 : : };
1621 : :
1622 : : static PyType_Spec takewhile_spec = {
1623 : : .name = "itertools.takewhile",
1624 : : .basicsize = sizeof(takewhileobject),
1625 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
1626 : : Py_TPFLAGS_IMMUTABLETYPE),
1627 : : .slots = takewhile_slots,
1628 : : };
1629 : :
1630 : :
1631 : : /* islice object *************************************************************/
1632 : :
1633 : : typedef struct {
1634 : : PyObject_HEAD
1635 : : PyObject *it;
1636 : : Py_ssize_t next;
1637 : : Py_ssize_t stop;
1638 : : Py_ssize_t step;
1639 : : Py_ssize_t cnt;
1640 : : } isliceobject;
1641 : :
1642 : : static PyObject *
1643 : 18 : islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
1644 : : {
1645 : : PyObject *seq;
1646 : 18 : Py_ssize_t start=0, stop=-1, step=1;
1647 : 18 : PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
1648 : : Py_ssize_t numargs;
1649 : : isliceobject *lz;
1650 : :
1651 : 18 : itertools_state *st = find_state_by_type(type);
1652 : 18 : PyTypeObject *islice_type = st->islice_type;
1653 [ - + - - : 18 : if ((type == islice_type || type->tp_init == islice_type->tp_init) &&
- + ]
1654 [ # # ]: 0 : !_PyArg_NoKeywords("islice", kwds))
1655 : 0 : return NULL;
1656 : :
1657 [ - + ]: 18 : if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
1658 : 0 : return NULL;
1659 : :
1660 : 18 : numargs = PyTuple_Size(args);
1661 [ - + ]: 18 : if (numargs == 2) {
1662 [ # # ]: 0 : if (a1 != Py_None) {
1663 : 0 : stop = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
1664 [ # # ]: 0 : if (stop == -1) {
1665 [ # # ]: 0 : if (PyErr_Occurred())
1666 : 0 : PyErr_Clear();
1667 : 0 : PyErr_SetString(PyExc_ValueError,
1668 : : "Stop argument for islice() must be None or "
1669 : : "an integer: 0 <= x <= sys.maxsize.");
1670 : 0 : return NULL;
1671 : : }
1672 : : }
1673 : : } else {
1674 [ + - ]: 18 : if (a1 != Py_None)
1675 : 18 : start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
1676 [ - + - - ]: 18 : if (start == -1 && PyErr_Occurred())
1677 : 0 : PyErr_Clear();
1678 [ - + ]: 18 : if (a2 != Py_None) {
1679 : 0 : stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
1680 [ # # ]: 0 : if (stop == -1) {
1681 [ # # ]: 0 : if (PyErr_Occurred())
1682 : 0 : PyErr_Clear();
1683 : 0 : PyErr_SetString(PyExc_ValueError,
1684 : : "Stop argument for islice() must be None or "
1685 : : "an integer: 0 <= x <= sys.maxsize.");
1686 : 0 : return NULL;
1687 : : }
1688 : : }
1689 : : }
1690 [ + - - + ]: 18 : if (start<0 || stop<-1) {
1691 : 0 : PyErr_SetString(PyExc_ValueError,
1692 : : "Indices for islice() must be None or "
1693 : : "an integer: 0 <= x <= sys.maxsize.");
1694 : 0 : return NULL;
1695 : : }
1696 : :
1697 [ - + ]: 18 : if (a3 != NULL) {
1698 [ # # ]: 0 : if (a3 != Py_None)
1699 : 0 : step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
1700 [ # # # # ]: 0 : if (step == -1 && PyErr_Occurred())
1701 : 0 : PyErr_Clear();
1702 : : }
1703 [ - + ]: 18 : if (step<1) {
1704 : 0 : PyErr_SetString(PyExc_ValueError,
1705 : : "Step for islice() must be a positive integer or None.");
1706 : 0 : return NULL;
1707 : : }
1708 : :
1709 : : /* Get iterator. */
1710 : 18 : it = PyObject_GetIter(seq);
1711 [ - + ]: 18 : if (it == NULL)
1712 : 0 : return NULL;
1713 : :
1714 : : /* create isliceobject structure */
1715 : 18 : lz = (isliceobject *)type->tp_alloc(type, 0);
1716 [ - + ]: 18 : if (lz == NULL) {
1717 : 0 : Py_DECREF(it);
1718 : 0 : return NULL;
1719 : : }
1720 : 18 : lz->it = it;
1721 : 18 : lz->next = start;
1722 : 18 : lz->stop = stop;
1723 : 18 : lz->step = step;
1724 : 18 : lz->cnt = 0L;
1725 : :
1726 : 18 : return (PyObject *)lz;
1727 : : }
1728 : :
1729 : : static void
1730 : 18 : islice_dealloc(isliceobject *lz)
1731 : : {
1732 : 18 : PyTypeObject *tp = Py_TYPE(lz);
1733 : 18 : PyObject_GC_UnTrack(lz);
1734 : 18 : Py_XDECREF(lz->it);
1735 : 18 : tp->tp_free(lz);
1736 : 18 : Py_DECREF(tp);
1737 : 18 : }
1738 : :
1739 : : static int
1740 : 0 : islice_traverse(isliceobject *lz, visitproc visit, void *arg)
1741 : : {
1742 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(lz));
1743 [ # # # # ]: 0 : Py_VISIT(lz->it);
1744 : 0 : return 0;
1745 : : }
1746 : :
1747 : : static PyObject *
1748 : 18 : islice_next(isliceobject *lz)
1749 : : {
1750 : : PyObject *item;
1751 : 18 : PyObject *it = lz->it;
1752 : 18 : Py_ssize_t stop = lz->stop;
1753 : : Py_ssize_t oldnext;
1754 : : PyObject *(*iternext)(PyObject *);
1755 : :
1756 [ - + ]: 18 : if (it == NULL)
1757 : 0 : return NULL;
1758 : :
1759 : 18 : iternext = *Py_TYPE(it)->tp_iternext;
1760 [ + + ]: 2879 : while (lz->cnt < lz->next) {
1761 : 2861 : item = iternext(it);
1762 [ - + ]: 2861 : if (item == NULL)
1763 : 0 : goto empty;
1764 : 2861 : Py_DECREF(item);
1765 : 2861 : lz->cnt++;
1766 : : }
1767 [ - + - - ]: 18 : if (stop != -1 && lz->cnt >= stop)
1768 : 0 : goto empty;
1769 : 18 : item = iternext(it);
1770 [ - + ]: 18 : if (item == NULL)
1771 : 0 : goto empty;
1772 : 18 : lz->cnt++;
1773 : 18 : oldnext = lz->next;
1774 : : /* The (size_t) cast below avoids the danger of undefined
1775 : : behaviour from signed integer overflow. */
1776 : 18 : lz->next += (size_t)lz->step;
1777 [ + - - + : 18 : if (lz->next < oldnext || (stop != -1 && lz->next > stop))
- - ]
1778 : 0 : lz->next = stop;
1779 : 18 : return item;
1780 : :
1781 : 0 : empty:
1782 [ # # ]: 0 : Py_CLEAR(lz->it);
1783 : 0 : return NULL;
1784 : : }
1785 : :
1786 : : static PyObject *
1787 : 0 : islice_reduce(isliceobject *lz, PyObject *Py_UNUSED(ignored))
1788 : : {
1789 : : /* When unpickled, generate a new object with the same bounds,
1790 : : * then 'setstate' with the next and count
1791 : : */
1792 : : PyObject *stop;
1793 : :
1794 [ # # ]: 0 : if (lz->it == NULL) {
1795 : : PyObject *empty_list;
1796 : : PyObject *empty_it;
1797 : 0 : empty_list = PyList_New(0);
1798 [ # # ]: 0 : if (empty_list == NULL)
1799 : 0 : return NULL;
1800 : 0 : empty_it = PyObject_GetIter(empty_list);
1801 : 0 : Py_DECREF(empty_list);
1802 [ # # ]: 0 : if (empty_it == NULL)
1803 : 0 : return NULL;
1804 : 0 : return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
1805 : : }
1806 [ # # ]: 0 : if (lz->stop == -1) {
1807 : 0 : stop = Py_NewRef(Py_None);
1808 : : } else {
1809 : 0 : stop = PyLong_FromSsize_t(lz->stop);
1810 [ # # ]: 0 : if (stop == NULL)
1811 : 0 : return NULL;
1812 : : }
1813 : 0 : return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
1814 : : lz->it, lz->next, stop, lz->step,
1815 : : lz->cnt);
1816 : : }
1817 : :
1818 : : static PyObject *
1819 : 0 : islice_setstate(isliceobject *lz, PyObject *state)
1820 : : {
1821 : 0 : Py_ssize_t cnt = PyLong_AsSsize_t(state);
1822 : :
1823 [ # # # # ]: 0 : if (cnt == -1 && PyErr_Occurred())
1824 : 0 : return NULL;
1825 : 0 : lz->cnt = cnt;
1826 : 0 : Py_RETURN_NONE;
1827 : : }
1828 : :
1829 : : static PyMethodDef islice_methods[] = {
1830 : : {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
1831 : : reduce_doc},
1832 : : {"__setstate__", (PyCFunction)islice_setstate, METH_O,
1833 : : setstate_doc},
1834 : : {NULL, NULL} /* sentinel */
1835 : : };
1836 : :
1837 : : PyDoc_STRVAR(islice_doc,
1838 : : "islice(iterable, stop) --> islice object\n\
1839 : : islice(iterable, start, stop[, step]) --> islice object\n\
1840 : : \n\
1841 : : Return an iterator whose next() method returns selected values from an\n\
1842 : : iterable. If start is specified, will skip all preceding elements;\n\
1843 : : otherwise, start defaults to zero. Step defaults to one. If\n\
1844 : : specified as another value, step determines how many values are\n\
1845 : : skipped between successive calls. Works like a slice() on a list\n\
1846 : : but returns an iterator.");
1847 : :
1848 : : static PyType_Slot islice_slots[] = {
1849 : : {Py_tp_dealloc, islice_dealloc},
1850 : : {Py_tp_getattro, PyObject_GenericGetAttr},
1851 : : {Py_tp_doc, (void *)islice_doc},
1852 : : {Py_tp_traverse, islice_traverse},
1853 : : {Py_tp_iter, PyObject_SelfIter},
1854 : : {Py_tp_iternext, islice_next},
1855 : : {Py_tp_methods, islice_methods},
1856 : : {Py_tp_new, islice_new},
1857 : : {Py_tp_free, PyObject_GC_Del},
1858 : : {0, NULL},
1859 : : };
1860 : :
1861 : : static PyType_Spec islice_spec = {
1862 : : .name = "itertools.islice",
1863 : : .basicsize = sizeof(isliceobject),
1864 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
1865 : : Py_TPFLAGS_IMMUTABLETYPE),
1866 : : .slots = islice_slots,
1867 : : };
1868 : :
1869 : :
1870 : : /* starmap object ************************************************************/
1871 : :
1872 : : typedef struct {
1873 : : PyObject_HEAD
1874 : : PyObject *func;
1875 : : PyObject *it;
1876 : : } starmapobject;
1877 : :
1878 : : /*[clinic input]
1879 : : @classmethod
1880 : : itertools.starmap.__new__
1881 : : function as func: object
1882 : : iterable as seq: object
1883 : : /
1884 : : Return an iterator whose values are returned from the function evaluated with an argument tuple taken from the given sequence.
1885 : : [clinic start generated code]*/
1886 : :
1887 : : static PyObject *
1888 : 0 : itertools_starmap_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
1889 : : /*[clinic end generated code: output=79eeb81d452c6e8d input=844766df6a0d4dad]*/
1890 : : {
1891 : : PyObject *it;
1892 : : starmapobject *lz;
1893 : :
1894 : : /* Get iterator. */
1895 : 0 : it = PyObject_GetIter(seq);
1896 [ # # ]: 0 : if (it == NULL)
1897 : 0 : return NULL;
1898 : :
1899 : : /* create starmapobject structure */
1900 : 0 : lz = (starmapobject *)type->tp_alloc(type, 0);
1901 [ # # ]: 0 : if (lz == NULL) {
1902 : 0 : Py_DECREF(it);
1903 : 0 : return NULL;
1904 : : }
1905 : 0 : lz->func = Py_NewRef(func);
1906 : 0 : lz->it = it;
1907 : :
1908 : 0 : return (PyObject *)lz;
1909 : : }
1910 : :
1911 : : static void
1912 : 0 : starmap_dealloc(starmapobject *lz)
1913 : : {
1914 : 0 : PyTypeObject *tp = Py_TYPE(lz);
1915 : 0 : PyObject_GC_UnTrack(lz);
1916 : 0 : Py_XDECREF(lz->func);
1917 : 0 : Py_XDECREF(lz->it);
1918 : 0 : tp->tp_free(lz);
1919 : 0 : Py_DECREF(tp);
1920 : 0 : }
1921 : :
1922 : : static int
1923 : 0 : starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
1924 : : {
1925 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(lz));
1926 [ # # # # ]: 0 : Py_VISIT(lz->it);
1927 [ # # # # ]: 0 : Py_VISIT(lz->func);
1928 : 0 : return 0;
1929 : : }
1930 : :
1931 : : static PyObject *
1932 : 0 : starmap_next(starmapobject *lz)
1933 : : {
1934 : : PyObject *args;
1935 : : PyObject *result;
1936 : 0 : PyObject *it = lz->it;
1937 : :
1938 : 0 : args = (*Py_TYPE(it)->tp_iternext)(it);
1939 [ # # ]: 0 : if (args == NULL)
1940 : 0 : return NULL;
1941 [ # # ]: 0 : if (!PyTuple_CheckExact(args)) {
1942 : 0 : PyObject *newargs = PySequence_Tuple(args);
1943 : 0 : Py_DECREF(args);
1944 [ # # ]: 0 : if (newargs == NULL)
1945 : 0 : return NULL;
1946 : 0 : args = newargs;
1947 : : }
1948 : 0 : result = PyObject_Call(lz->func, args, NULL);
1949 : 0 : Py_DECREF(args);
1950 : 0 : return result;
1951 : : }
1952 : :
1953 : : static PyObject *
1954 : 0 : starmap_reduce(starmapobject *lz, PyObject *Py_UNUSED(ignored))
1955 : : {
1956 : : /* Just pickle the iterator */
1957 : 0 : return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
1958 : : }
1959 : :
1960 : : static PyMethodDef starmap_methods[] = {
1961 : : {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
1962 : : reduce_doc},
1963 : : {NULL, NULL} /* sentinel */
1964 : : };
1965 : :
1966 : : static PyType_Slot starmap_slots[] = {
1967 : : {Py_tp_dealloc, starmap_dealloc},
1968 : : {Py_tp_getattro, PyObject_GenericGetAttr},
1969 : : {Py_tp_doc, (void *)itertools_starmap__doc__},
1970 : : {Py_tp_traverse, starmap_traverse},
1971 : : {Py_tp_iter, PyObject_SelfIter},
1972 : : {Py_tp_iternext, starmap_next},
1973 : : {Py_tp_methods, starmap_methods},
1974 : : {Py_tp_new, itertools_starmap},
1975 : : {Py_tp_free, PyObject_GC_Del},
1976 : : {0, NULL},
1977 : : };
1978 : :
1979 : : static PyType_Spec starmap_spec = {
1980 : : .name = "itertools.starmap",
1981 : : .basicsize = sizeof(starmapobject),
1982 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
1983 : : Py_TPFLAGS_IMMUTABLETYPE),
1984 : : .slots = starmap_slots,
1985 : : };
1986 : :
1987 : :
1988 : : /* chain object **************************************************************/
1989 : :
1990 : : typedef struct {
1991 : : PyObject_HEAD
1992 : : PyObject *source; /* Iterator over input iterables */
1993 : : PyObject *active; /* Currently running input iterator */
1994 : : } chainobject;
1995 : :
1996 : : static PyObject *
1997 : 0 : chain_new_internal(PyTypeObject *type, PyObject *source)
1998 : : {
1999 : : chainobject *lz;
2000 : :
2001 : 0 : lz = (chainobject *)type->tp_alloc(type, 0);
2002 [ # # ]: 0 : if (lz == NULL) {
2003 : 0 : Py_DECREF(source);
2004 : 0 : return NULL;
2005 : : }
2006 : :
2007 : 0 : lz->source = source;
2008 : 0 : lz->active = NULL;
2009 : 0 : return (PyObject *)lz;
2010 : : }
2011 : :
2012 : : static PyObject *
2013 : 0 : chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2014 : : {
2015 : : PyObject *source;
2016 : :
2017 : 0 : itertools_state *state = find_state_by_type(type);
2018 : 0 : PyTypeObject *chain_type = state->chain_type;
2019 [ # # # # : 0 : if ((type == chain_type || type->tp_init == chain_type->tp_init) &&
# # ]
2020 [ # # ]: 0 : !_PyArg_NoKeywords("chain", kwds))
2021 : 0 : return NULL;
2022 : :
2023 : 0 : source = PyObject_GetIter(args);
2024 [ # # ]: 0 : if (source == NULL)
2025 : 0 : return NULL;
2026 : :
2027 : 0 : return chain_new_internal(type, source);
2028 : : }
2029 : :
2030 : : /*[clinic input]
2031 : : @classmethod
2032 : : itertools.chain.from_iterable
2033 : : iterable as arg: object
2034 : : /
2035 : : Alternative chain() constructor taking a single iterable argument that evaluates lazily.
2036 : : [clinic start generated code]*/
2037 : :
2038 : : static PyObject *
2039 : 0 : itertools_chain_from_iterable(PyTypeObject *type, PyObject *arg)
2040 : : /*[clinic end generated code: output=667ae7a7f7b68654 input=72c39e3a2ca3be85]*/
2041 : : {
2042 : : PyObject *source;
2043 : :
2044 : 0 : source = PyObject_GetIter(arg);
2045 [ # # ]: 0 : if (source == NULL)
2046 : 0 : return NULL;
2047 : :
2048 : 0 : return chain_new_internal(type, source);
2049 : : }
2050 : :
2051 : : static void
2052 : 0 : chain_dealloc(chainobject *lz)
2053 : : {
2054 : 0 : PyTypeObject *tp = Py_TYPE(lz);
2055 : 0 : PyObject_GC_UnTrack(lz);
2056 : 0 : Py_XDECREF(lz->active);
2057 : 0 : Py_XDECREF(lz->source);
2058 : 0 : tp->tp_free(lz);
2059 : 0 : Py_DECREF(tp);
2060 : 0 : }
2061 : :
2062 : : static int
2063 : 0 : chain_traverse(chainobject *lz, visitproc visit, void *arg)
2064 : : {
2065 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(lz));
2066 [ # # # # ]: 0 : Py_VISIT(lz->source);
2067 [ # # # # ]: 0 : Py_VISIT(lz->active);
2068 : 0 : return 0;
2069 : : }
2070 : :
2071 : : static PyObject *
2072 : 0 : chain_next(chainobject *lz)
2073 : : {
2074 : : PyObject *item;
2075 : :
2076 : : /* lz->source is the iterator of iterables. If it's NULL, we've already
2077 : : * consumed them all. lz->active is the current iterator. If it's NULL,
2078 : : * we should grab a new one from lz->source. */
2079 [ # # ]: 0 : while (lz->source != NULL) {
2080 [ # # ]: 0 : if (lz->active == NULL) {
2081 : 0 : PyObject *iterable = PyIter_Next(lz->source);
2082 [ # # ]: 0 : if (iterable == NULL) {
2083 [ # # ]: 0 : Py_CLEAR(lz->source);
2084 : 0 : return NULL; /* no more input sources */
2085 : : }
2086 : 0 : lz->active = PyObject_GetIter(iterable);
2087 : 0 : Py_DECREF(iterable);
2088 [ # # ]: 0 : if (lz->active == NULL) {
2089 [ # # ]: 0 : Py_CLEAR(lz->source);
2090 : 0 : return NULL; /* input not iterable */
2091 : : }
2092 : : }
2093 : 0 : item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
2094 [ # # ]: 0 : if (item != NULL)
2095 : 0 : return item;
2096 [ # # ]: 0 : if (PyErr_Occurred()) {
2097 [ # # ]: 0 : if (PyErr_ExceptionMatches(PyExc_StopIteration))
2098 : 0 : PyErr_Clear();
2099 : : else
2100 : 0 : return NULL; /* input raised an exception */
2101 : : }
2102 : : /* lz->active is consumed, try with the next iterable. */
2103 [ # # ]: 0 : Py_CLEAR(lz->active);
2104 : : }
2105 : : /* Everything had been consumed already. */
2106 : 0 : return NULL;
2107 : : }
2108 : :
2109 : : static PyObject *
2110 : 0 : chain_reduce(chainobject *lz, PyObject *Py_UNUSED(ignored))
2111 : : {
2112 [ # # ]: 0 : if (lz->source) {
2113 : : /* we can't pickle function objects (itertools.from_iterable) so
2114 : : * we must use setstate to replace the iterable. One day we
2115 : : * will fix pickling of functions
2116 : : */
2117 [ # # ]: 0 : if (lz->active) {
2118 : 0 : return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
2119 : : } else {
2120 : 0 : return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
2121 : : }
2122 : : } else {
2123 : 0 : return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
2124 : : }
2125 : : return NULL;
2126 : : }
2127 : :
2128 : : static PyObject *
2129 : 0 : chain_setstate(chainobject *lz, PyObject *state)
2130 : : {
2131 : 0 : PyObject *source, *active=NULL;
2132 : :
2133 [ # # ]: 0 : if (!PyTuple_Check(state)) {
2134 : 0 : PyErr_SetString(PyExc_TypeError, "state is not a tuple");
2135 : 0 : return NULL;
2136 : : }
2137 [ # # ]: 0 : if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
2138 : 0 : return NULL;
2139 : : }
2140 [ # # # # : 0 : if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
# # ]
2141 : 0 : PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
2142 : 0 : return NULL;
2143 : : }
2144 : :
2145 : 0 : Py_INCREF(source);
2146 : 0 : Py_XSETREF(lz->source, source);
2147 : 0 : Py_XINCREF(active);
2148 : 0 : Py_XSETREF(lz->active, active);
2149 : 0 : Py_RETURN_NONE;
2150 : : }
2151 : :
2152 : : PyDoc_STRVAR(chain_doc,
2153 : : "chain(*iterables) --> chain object\n\
2154 : : \n\
2155 : : Return a chain object whose .__next__() method returns elements from the\n\
2156 : : first iterable until it is exhausted, then elements from the next\n\
2157 : : iterable, until all of the iterables are exhausted.");
2158 : :
2159 : : static PyMethodDef chain_methods[] = {
2160 : : ITERTOOLS_CHAIN_FROM_ITERABLE_METHODDEF
2161 : : {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
2162 : : reduce_doc},
2163 : : {"__setstate__", (PyCFunction)chain_setstate, METH_O,
2164 : : setstate_doc},
2165 : : {"__class_getitem__", Py_GenericAlias,
2166 : : METH_O|METH_CLASS, PyDoc_STR("See PEP 585")},
2167 : : {NULL, NULL} /* sentinel */
2168 : : };
2169 : :
2170 : : static PyType_Slot chain_slots[] = {
2171 : : {Py_tp_dealloc, chain_dealloc},
2172 : : {Py_tp_getattro, PyObject_GenericGetAttr},
2173 : : {Py_tp_doc, (void *)chain_doc},
2174 : : {Py_tp_traverse, chain_traverse},
2175 : : {Py_tp_iter, PyObject_SelfIter},
2176 : : {Py_tp_iternext, chain_next},
2177 : : {Py_tp_methods, chain_methods},
2178 : : {Py_tp_new, chain_new},
2179 : : {Py_tp_free, PyObject_GC_Del},
2180 : : {0, NULL},
2181 : : };
2182 : :
2183 : : static PyType_Spec chain_spec = {
2184 : : .name = "itertools.chain",
2185 : : .basicsize = sizeof(chainobject),
2186 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
2187 : : Py_TPFLAGS_IMMUTABLETYPE),
2188 : : .slots = chain_slots,
2189 : : };
2190 : :
2191 : :
2192 : : /* product object ************************************************************/
2193 : :
2194 : : typedef struct {
2195 : : PyObject_HEAD
2196 : : PyObject *pools; /* tuple of pool tuples */
2197 : : Py_ssize_t *indices; /* one index per pool */
2198 : : PyObject *result; /* most recently returned result tuple */
2199 : : int stopped; /* set to 1 when the iterator is exhausted */
2200 : : } productobject;
2201 : :
2202 : : static PyObject *
2203 : 416 : product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2204 : : {
2205 : : productobject *lz;
2206 : 416 : Py_ssize_t nargs, npools, repeat=1;
2207 : 416 : PyObject *pools = NULL;
2208 : 416 : Py_ssize_t *indices = NULL;
2209 : : Py_ssize_t i;
2210 : :
2211 [ + + ]: 416 : if (kwds != NULL) {
2212 : 344 : char *kwlist[] = {"repeat", 0};
2213 : 344 : PyObject *tmpargs = PyTuple_New(0);
2214 [ - + ]: 344 : if (tmpargs == NULL)
2215 : 0 : return NULL;
2216 [ - + ]: 344 : if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
2217 : : kwlist, &repeat)) {
2218 : 0 : Py_DECREF(tmpargs);
2219 : 0 : return NULL;
2220 : : }
2221 : 344 : Py_DECREF(tmpargs);
2222 [ - + ]: 344 : if (repeat < 0) {
2223 : 0 : PyErr_SetString(PyExc_ValueError,
2224 : : "repeat argument cannot be negative");
2225 : 0 : return NULL;
2226 : : }
2227 : : }
2228 : :
2229 : : assert(PyTuple_CheckExact(args));
2230 [ - + ]: 416 : if (repeat == 0) {
2231 : 0 : nargs = 0;
2232 : : } else {
2233 : 416 : nargs = PyTuple_GET_SIZE(args);
2234 [ - + ]: 416 : if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
2235 : 0 : PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
2236 : 0 : return NULL;
2237 : : }
2238 : : }
2239 : 416 : npools = nargs * repeat;
2240 : :
2241 [ + - ]: 416 : indices = PyMem_New(Py_ssize_t, npools);
2242 [ - + ]: 416 : if (indices == NULL) {
2243 : 0 : PyErr_NoMemory();
2244 : 0 : goto error;
2245 : : }
2246 : :
2247 : 416 : pools = PyTuple_New(npools);
2248 [ - + ]: 416 : if (pools == NULL)
2249 : 0 : goto error;
2250 : :
2251 [ + + ]: 868 : for (i=0; i < nargs ; ++i) {
2252 : 452 : PyObject *item = PyTuple_GET_ITEM(args, i);
2253 : 452 : PyObject *pool = PySequence_Tuple(item);
2254 [ - + ]: 452 : if (pool == NULL)
2255 : 0 : goto error;
2256 : 452 : PyTuple_SET_ITEM(pools, i, pool);
2257 : 452 : indices[i] = 0;
2258 : : }
2259 [ + + ]: 1104 : for ( ; i < npools; ++i) {
2260 : 688 : PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
2261 : 688 : Py_INCREF(pool);
2262 : 688 : PyTuple_SET_ITEM(pools, i, pool);
2263 : 688 : indices[i] = 0;
2264 : : }
2265 : :
2266 : : /* create productobject structure */
2267 : 416 : lz = (productobject *)type->tp_alloc(type, 0);
2268 [ - + ]: 416 : if (lz == NULL)
2269 : 0 : goto error;
2270 : :
2271 : 416 : lz->pools = pools;
2272 : 416 : lz->indices = indices;
2273 : 416 : lz->result = NULL;
2274 : 416 : lz->stopped = 0;
2275 : :
2276 : 416 : return (PyObject *)lz;
2277 : :
2278 : 0 : error:
2279 [ # # ]: 0 : if (indices != NULL)
2280 : 0 : PyMem_Free(indices);
2281 : 0 : Py_XDECREF(pools);
2282 : 0 : return NULL;
2283 : : }
2284 : :
2285 : : static void
2286 : 416 : product_dealloc(productobject *lz)
2287 : : {
2288 : 416 : PyTypeObject *tp = Py_TYPE(lz);
2289 : 416 : PyObject_GC_UnTrack(lz);
2290 : 416 : Py_XDECREF(lz->pools);
2291 : 416 : Py_XDECREF(lz->result);
2292 [ + - ]: 416 : if (lz->indices != NULL)
2293 : 416 : PyMem_Free(lz->indices);
2294 : 416 : tp->tp_free(lz);
2295 : 416 : Py_DECREF(tp);
2296 : 416 : }
2297 : :
2298 : : static PyObject *
2299 : 0 : product_sizeof(productobject *lz, void *unused)
2300 : : {
2301 : 0 : size_t res = _PyObject_SIZE(Py_TYPE(lz));
2302 : 0 : res += (size_t)PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
2303 : 0 : return PyLong_FromSize_t(res);
2304 : : }
2305 : :
2306 : : PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
2307 : :
2308 : : static int
2309 : 0 : product_traverse(productobject *lz, visitproc visit, void *arg)
2310 : : {
2311 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(lz));
2312 [ # # # # ]: 0 : Py_VISIT(lz->pools);
2313 [ # # # # ]: 0 : Py_VISIT(lz->result);
2314 : 0 : return 0;
2315 : : }
2316 : :
2317 : : static PyObject *
2318 : 118624 : product_next(productobject *lz)
2319 : : {
2320 : : PyObject *pool;
2321 : : PyObject *elem;
2322 : : PyObject *oldelem;
2323 : 118624 : PyObject *pools = lz->pools;
2324 : 118624 : PyObject *result = lz->result;
2325 : 118624 : Py_ssize_t npools = PyTuple_GET_SIZE(pools);
2326 : : Py_ssize_t i;
2327 : :
2328 [ - + ]: 118624 : if (lz->stopped)
2329 : 0 : return NULL;
2330 : :
2331 [ + + ]: 118624 : if (result == NULL) {
2332 : : /* On the first pass, return an initial tuple filled with the
2333 : : first element from each pool. */
2334 : 416 : result = PyTuple_New(npools);
2335 [ - + ]: 416 : if (result == NULL)
2336 : 0 : goto empty;
2337 : 416 : lz->result = result;
2338 [ + + ]: 1556 : for (i=0; i < npools; i++) {
2339 : 1140 : pool = PyTuple_GET_ITEM(pools, i);
2340 [ - + ]: 1140 : if (PyTuple_GET_SIZE(pool) == 0)
2341 : 0 : goto empty;
2342 : 1140 : elem = PyTuple_GET_ITEM(pool, 0);
2343 : 1140 : Py_INCREF(elem);
2344 : 1140 : PyTuple_SET_ITEM(result, i, elem);
2345 : : }
2346 : : } else {
2347 : 118208 : Py_ssize_t *indices = lz->indices;
2348 : :
2349 : : /* Copy the previous result tuple or re-use it if available */
2350 [ + - ]: 118208 : if (Py_REFCNT(result) > 1) {
2351 : 118208 : PyObject *old_result = result;
2352 : 118208 : result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), npools);
2353 [ - + ]: 118208 : if (result == NULL)
2354 : 0 : goto empty;
2355 : 118208 : lz->result = result;
2356 : 118208 : Py_DECREF(old_result);
2357 : : }
2358 : : // bpo-42536: The GC may have untracked this result tuple. Since we're
2359 : : // recycling it, make sure it's tracked again:
2360 [ # # ]: 0 : else if (!_PyObject_GC_IS_TRACKED(result)) {
2361 : 0 : _PyObject_GC_TRACK(result);
2362 : : }
2363 : : /* Now, we've got the only copy so we can update it in-place */
2364 : : assert (npools==0 || Py_REFCNT(result) == 1);
2365 : :
2366 : : /* Update the pool indices right-to-left. Only advance to the
2367 : : next pool when the previous one rolls-over */
2368 [ + + ]: 137960 : for (i=npools-1 ; i >= 0 ; i--) {
2369 : 137544 : pool = PyTuple_GET_ITEM(pools, i);
2370 : 137544 : indices[i]++;
2371 [ + + ]: 137544 : if (indices[i] == PyTuple_GET_SIZE(pool)) {
2372 : : /* Roll-over and advance to next pool */
2373 : 19752 : indices[i] = 0;
2374 : 19752 : elem = PyTuple_GET_ITEM(pool, 0);
2375 : 19752 : Py_INCREF(elem);
2376 : 19752 : oldelem = PyTuple_GET_ITEM(result, i);
2377 : 19752 : PyTuple_SET_ITEM(result, i, elem);
2378 : 19752 : Py_DECREF(oldelem);
2379 : : } else {
2380 : : /* No rollover. Just increment and stop here. */
2381 : 117792 : elem = PyTuple_GET_ITEM(pool, indices[i]);
2382 : 117792 : Py_INCREF(elem);
2383 : 117792 : oldelem = PyTuple_GET_ITEM(result, i);
2384 : 117792 : PyTuple_SET_ITEM(result, i, elem);
2385 : 117792 : Py_DECREF(oldelem);
2386 : 117792 : break;
2387 : : }
2388 : : }
2389 : :
2390 : : /* If i is negative, then the indices have all rolled-over
2391 : : and we're done. */
2392 [ + + ]: 118208 : if (i < 0)
2393 : 416 : goto empty;
2394 : : }
2395 : :
2396 : 118208 : return Py_NewRef(result);
2397 : :
2398 : 416 : empty:
2399 : 416 : lz->stopped = 1;
2400 : 416 : return NULL;
2401 : : }
2402 : :
2403 : : static PyObject *
2404 : 0 : product_reduce(productobject *lz, PyObject *Py_UNUSED(ignored))
2405 : : {
2406 [ # # ]: 0 : if (lz->stopped) {
2407 : 0 : return Py_BuildValue("O(())", Py_TYPE(lz));
2408 [ # # ]: 0 : } else if (lz->result == NULL) {
2409 : 0 : return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
2410 : : } else {
2411 : : PyObject *indices;
2412 : : Py_ssize_t n, i;
2413 : :
2414 : : /* we must pickle the indices use them for setstate, and
2415 : : * additionally indicate that the iterator has started
2416 : : */
2417 : 0 : n = PyTuple_GET_SIZE(lz->pools);
2418 : 0 : indices = PyTuple_New(n);
2419 [ # # ]: 0 : if (indices == NULL)
2420 : 0 : return NULL;
2421 [ # # ]: 0 : for (i=0; i<n; i++){
2422 : 0 : PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2423 [ # # ]: 0 : if (!index) {
2424 : 0 : Py_DECREF(indices);
2425 : 0 : return NULL;
2426 : : }
2427 : 0 : PyTuple_SET_ITEM(indices, i, index);
2428 : : }
2429 : 0 : return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
2430 : : }
2431 : : }
2432 : :
2433 : : static PyObject *
2434 : 0 : product_setstate(productobject *lz, PyObject *state)
2435 : : {
2436 : : PyObject *result;
2437 : : Py_ssize_t n, i;
2438 : :
2439 : 0 : n = PyTuple_GET_SIZE(lz->pools);
2440 [ # # # # ]: 0 : if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
2441 : 0 : PyErr_SetString(PyExc_ValueError, "invalid arguments");
2442 : 0 : return NULL;
2443 : : }
2444 [ # # ]: 0 : for (i=0; i<n; i++)
2445 : : {
2446 : 0 : PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2447 : 0 : Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2448 : : PyObject* pool;
2449 : : Py_ssize_t poolsize;
2450 [ # # # # ]: 0 : if (index < 0 && PyErr_Occurred())
2451 : 0 : return NULL; /* not an integer */
2452 : 0 : pool = PyTuple_GET_ITEM(lz->pools, i);
2453 : 0 : poolsize = PyTuple_GET_SIZE(pool);
2454 [ # # ]: 0 : if (poolsize == 0) {
2455 : 0 : lz->stopped = 1;
2456 : 0 : Py_RETURN_NONE;
2457 : : }
2458 : : /* clamp the index */
2459 [ # # ]: 0 : if (index < 0)
2460 : 0 : index = 0;
2461 [ # # ]: 0 : else if (index > poolsize-1)
2462 : 0 : index = poolsize-1;
2463 : 0 : lz->indices[i] = index;
2464 : : }
2465 : :
2466 : 0 : result = PyTuple_New(n);
2467 [ # # ]: 0 : if (!result)
2468 : 0 : return NULL;
2469 [ # # ]: 0 : for (i=0; i<n; i++) {
2470 : 0 : PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
2471 : 0 : PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
2472 : 0 : Py_INCREF(element);
2473 : 0 : PyTuple_SET_ITEM(result, i, element);
2474 : : }
2475 : 0 : Py_XSETREF(lz->result, result);
2476 : 0 : Py_RETURN_NONE;
2477 : : }
2478 : :
2479 : : static PyMethodDef product_methods[] = {
2480 : : {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
2481 : : reduce_doc},
2482 : : {"__setstate__", (PyCFunction)product_setstate, METH_O,
2483 : : setstate_doc},
2484 : : {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
2485 : : sizeof_doc},
2486 : : {NULL, NULL} /* sentinel */
2487 : : };
2488 : :
2489 : : PyDoc_STRVAR(product_doc,
2490 : : "product(*iterables, repeat=1) --> product object\n\
2491 : : \n\
2492 : : Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
2493 : : For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
2494 : : The leftmost iterators are in the outermost for-loop, so the output tuples\n\
2495 : : cycle in a manner similar to an odometer (with the rightmost element changing\n\
2496 : : on every iteration).\n\n\
2497 : : To compute the product of an iterable with itself, specify the number\n\
2498 : : of repetitions with the optional repeat keyword argument. For example,\n\
2499 : : product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
2500 : : product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
2501 : : product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
2502 : :
2503 : : static PyType_Slot product_slots[] = {
2504 : : {Py_tp_dealloc, product_dealloc},
2505 : : {Py_tp_getattro, PyObject_GenericGetAttr},
2506 : : {Py_tp_doc, (void *)product_doc},
2507 : : {Py_tp_traverse, product_traverse},
2508 : : {Py_tp_iter, PyObject_SelfIter},
2509 : : {Py_tp_iternext, product_next},
2510 : : {Py_tp_methods, product_methods},
2511 : : {Py_tp_new, product_new},
2512 : : {Py_tp_free, PyObject_GC_Del},
2513 : : {0, NULL},
2514 : : };
2515 : :
2516 : : static PyType_Spec product_spec = {
2517 : : .name = "itertools.product",
2518 : : .basicsize = sizeof(productobject),
2519 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
2520 : : Py_TPFLAGS_IMMUTABLETYPE),
2521 : : .slots = product_slots,
2522 : : };
2523 : :
2524 : :
2525 : : /* combinations object *******************************************************/
2526 : :
2527 : : typedef struct {
2528 : : PyObject_HEAD
2529 : : PyObject *pool; /* input converted to a tuple */
2530 : : Py_ssize_t *indices; /* one index per result element */
2531 : : PyObject *result; /* most recently returned result tuple */
2532 : : Py_ssize_t r; /* size of result tuple */
2533 : : int stopped; /* set to 1 when the iterator is exhausted */
2534 : : } combinationsobject;
2535 : :
2536 : :
2537 : : /*[clinic input]
2538 : : @classmethod
2539 : : itertools.combinations.__new__
2540 : : iterable: object
2541 : : r: Py_ssize_t
2542 : : Return successive r-length combinations of elements in the iterable.
2543 : :
2544 : : combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)
2545 : : [clinic start generated code]*/
2546 : :
2547 : : static PyObject *
2548 : 0 : itertools_combinations_impl(PyTypeObject *type, PyObject *iterable,
2549 : : Py_ssize_t r)
2550 : : /*[clinic end generated code: output=87a689b39c40039c input=06bede09e3da20f8]*/
2551 : : {
2552 : : combinationsobject *co;
2553 : : Py_ssize_t n;
2554 : 0 : PyObject *pool = NULL;
2555 : 0 : Py_ssize_t *indices = NULL;
2556 : : Py_ssize_t i;
2557 : :
2558 : 0 : pool = PySequence_Tuple(iterable);
2559 [ # # ]: 0 : if (pool == NULL)
2560 : 0 : goto error;
2561 : 0 : n = PyTuple_GET_SIZE(pool);
2562 [ # # ]: 0 : if (r < 0) {
2563 : 0 : PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2564 : 0 : goto error;
2565 : : }
2566 : :
2567 [ # # ]: 0 : indices = PyMem_New(Py_ssize_t, r);
2568 [ # # ]: 0 : if (indices == NULL) {
2569 : 0 : PyErr_NoMemory();
2570 : 0 : goto error;
2571 : : }
2572 : :
2573 [ # # ]: 0 : for (i=0 ; i<r ; i++)
2574 : 0 : indices[i] = i;
2575 : :
2576 : : /* create combinationsobject structure */
2577 : 0 : co = (combinationsobject *)type->tp_alloc(type, 0);
2578 [ # # ]: 0 : if (co == NULL)
2579 : 0 : goto error;
2580 : :
2581 : 0 : co->pool = pool;
2582 : 0 : co->indices = indices;
2583 : 0 : co->result = NULL;
2584 : 0 : co->r = r;
2585 : 0 : co->stopped = r > n ? 1 : 0;
2586 : :
2587 : 0 : return (PyObject *)co;
2588 : :
2589 : 0 : error:
2590 [ # # ]: 0 : if (indices != NULL)
2591 : 0 : PyMem_Free(indices);
2592 : 0 : Py_XDECREF(pool);
2593 : 0 : return NULL;
2594 : : }
2595 : :
2596 : : static void
2597 : 0 : combinations_dealloc(combinationsobject *co)
2598 : : {
2599 : 0 : PyTypeObject *tp = Py_TYPE(co);
2600 : 0 : PyObject_GC_UnTrack(co);
2601 : 0 : Py_XDECREF(co->pool);
2602 : 0 : Py_XDECREF(co->result);
2603 [ # # ]: 0 : if (co->indices != NULL)
2604 : 0 : PyMem_Free(co->indices);
2605 : 0 : tp->tp_free(co);
2606 : 0 : Py_DECREF(tp);
2607 : 0 : }
2608 : :
2609 : : static PyObject *
2610 : 0 : combinations_sizeof(combinationsobject *co, void *unused)
2611 : : {
2612 : 0 : size_t res = _PyObject_SIZE(Py_TYPE(co));
2613 : 0 : res += (size_t)co->r * sizeof(Py_ssize_t);
2614 : 0 : return PyLong_FromSize_t(res);
2615 : : }
2616 : :
2617 : : static int
2618 : 0 : combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
2619 : : {
2620 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(co));
2621 [ # # # # ]: 0 : Py_VISIT(co->pool);
2622 [ # # # # ]: 0 : Py_VISIT(co->result);
2623 : 0 : return 0;
2624 : : }
2625 : :
2626 : : static PyObject *
2627 : 0 : combinations_next(combinationsobject *co)
2628 : : {
2629 : : PyObject *elem;
2630 : : PyObject *oldelem;
2631 : 0 : PyObject *pool = co->pool;
2632 : 0 : Py_ssize_t *indices = co->indices;
2633 : 0 : PyObject *result = co->result;
2634 : 0 : Py_ssize_t n = PyTuple_GET_SIZE(pool);
2635 : 0 : Py_ssize_t r = co->r;
2636 : : Py_ssize_t i, j, index;
2637 : :
2638 [ # # ]: 0 : if (co->stopped)
2639 : 0 : return NULL;
2640 : :
2641 [ # # ]: 0 : if (result == NULL) {
2642 : : /* On the first pass, initialize result tuple using the indices */
2643 : 0 : result = PyTuple_New(r);
2644 [ # # ]: 0 : if (result == NULL)
2645 : 0 : goto empty;
2646 : 0 : co->result = result;
2647 [ # # ]: 0 : for (i=0; i<r ; i++) {
2648 : 0 : index = indices[i];
2649 : 0 : elem = PyTuple_GET_ITEM(pool, index);
2650 : 0 : Py_INCREF(elem);
2651 : 0 : PyTuple_SET_ITEM(result, i, elem);
2652 : : }
2653 : : } else {
2654 : : /* Copy the previous result tuple or re-use it if available */
2655 [ # # ]: 0 : if (Py_REFCNT(result) > 1) {
2656 : 0 : PyObject *old_result = result;
2657 : 0 : result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
2658 [ # # ]: 0 : if (result == NULL)
2659 : 0 : goto empty;
2660 : 0 : co->result = result;
2661 : 0 : Py_DECREF(old_result);
2662 : : }
2663 : : // bpo-42536: The GC may have untracked this result tuple. Since we're
2664 : : // recycling it, make sure it's tracked again:
2665 [ # # ]: 0 : else if (!_PyObject_GC_IS_TRACKED(result)) {
2666 : 0 : _PyObject_GC_TRACK(result);
2667 : : }
2668 : : /* Now, we've got the only copy so we can update it in-place
2669 : : * CPython's empty tuple is a singleton and cached in
2670 : : * PyTuple's freelist.
2671 : : */
2672 : : assert(r == 0 || Py_REFCNT(result) == 1);
2673 : :
2674 : : /* Scan indices right-to-left until finding one that is not
2675 : : at its maximum (i + n - r). */
2676 [ # # # # ]: 0 : for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
2677 : : ;
2678 : :
2679 : : /* If i is negative, then the indices are all at
2680 : : their maximum value and we're done. */
2681 [ # # ]: 0 : if (i < 0)
2682 : 0 : goto empty;
2683 : :
2684 : : /* Increment the current index which we know is not at its
2685 : : maximum. Then move back to the right setting each index
2686 : : to its lowest possible value (one higher than the index
2687 : : to its left -- this maintains the sort order invariant). */
2688 : 0 : indices[i]++;
2689 [ # # ]: 0 : for (j=i+1 ; j<r ; j++)
2690 : 0 : indices[j] = indices[j-1] + 1;
2691 : :
2692 : : /* Update the result tuple for the new indices
2693 : : starting with i, the leftmost index that changed */
2694 [ # # ]: 0 : for ( ; i<r ; i++) {
2695 : 0 : index = indices[i];
2696 : 0 : elem = PyTuple_GET_ITEM(pool, index);
2697 : 0 : Py_INCREF(elem);
2698 : 0 : oldelem = PyTuple_GET_ITEM(result, i);
2699 : 0 : PyTuple_SET_ITEM(result, i, elem);
2700 : 0 : Py_DECREF(oldelem);
2701 : : }
2702 : : }
2703 : :
2704 : 0 : return Py_NewRef(result);
2705 : :
2706 : 0 : empty:
2707 : 0 : co->stopped = 1;
2708 : 0 : return NULL;
2709 : : }
2710 : :
2711 : : static PyObject *
2712 : 0 : combinations_reduce(combinationsobject *lz, PyObject *Py_UNUSED(ignored))
2713 : : {
2714 [ # # ]: 0 : if (lz->result == NULL) {
2715 : 0 : return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
2716 [ # # ]: 0 : } else if (lz->stopped) {
2717 : 0 : return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
2718 : : } else {
2719 : : PyObject *indices;
2720 : : Py_ssize_t i;
2721 : :
2722 : : /* we must pickle the indices and use them for setstate */
2723 : 0 : indices = PyTuple_New(lz->r);
2724 [ # # ]: 0 : if (!indices)
2725 : 0 : return NULL;
2726 [ # # ]: 0 : for (i=0; i<lz->r; i++)
2727 : : {
2728 : 0 : PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
2729 [ # # ]: 0 : if (!index) {
2730 : 0 : Py_DECREF(indices);
2731 : 0 : return NULL;
2732 : : }
2733 : 0 : PyTuple_SET_ITEM(indices, i, index);
2734 : : }
2735 : :
2736 : 0 : return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
2737 : : }
2738 : : }
2739 : :
2740 : : static PyObject *
2741 : 0 : combinations_setstate(combinationsobject *lz, PyObject *state)
2742 : : {
2743 : : PyObject *result;
2744 : : Py_ssize_t i;
2745 : 0 : Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
2746 : :
2747 [ # # # # ]: 0 : if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
2748 : 0 : PyErr_SetString(PyExc_ValueError, "invalid arguments");
2749 : 0 : return NULL;
2750 : : }
2751 : :
2752 [ # # ]: 0 : for (i=0; i<lz->r; i++) {
2753 : : Py_ssize_t max;
2754 : 0 : PyObject* indexObject = PyTuple_GET_ITEM(state, i);
2755 : 0 : Py_ssize_t index = PyLong_AsSsize_t(indexObject);
2756 : :
2757 [ # # # # ]: 0 : if (index == -1 && PyErr_Occurred())
2758 : 0 : return NULL; /* not an integer */
2759 : 0 : max = i + n - lz->r;
2760 : : /* clamp the index (beware of negative max) */
2761 [ # # ]: 0 : if (index > max)
2762 : 0 : index = max;
2763 [ # # ]: 0 : if (index < 0)
2764 : 0 : index = 0;
2765 : 0 : lz->indices[i] = index;
2766 : : }
2767 : :
2768 : 0 : result = PyTuple_New(lz->r);
2769 [ # # ]: 0 : if (result == NULL)
2770 : 0 : return NULL;
2771 [ # # ]: 0 : for (i=0; i<lz->r; i++) {
2772 : 0 : PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
2773 : 0 : Py_INCREF(element);
2774 : 0 : PyTuple_SET_ITEM(result, i, element);
2775 : : }
2776 : :
2777 : 0 : Py_XSETREF(lz->result, result);
2778 : 0 : Py_RETURN_NONE;
2779 : : }
2780 : :
2781 : : static PyMethodDef combinations_methods[] = {
2782 : : {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
2783 : : reduce_doc},
2784 : : {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
2785 : : setstate_doc},
2786 : : {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
2787 : : sizeof_doc},
2788 : : {NULL, NULL} /* sentinel */
2789 : : };
2790 : :
2791 : : static PyType_Slot combinations_slots[] = {
2792 : : {Py_tp_dealloc, combinations_dealloc},
2793 : : {Py_tp_getattro, PyObject_GenericGetAttr},
2794 : : {Py_tp_doc, (void *)itertools_combinations__doc__},
2795 : : {Py_tp_traverse, combinations_traverse},
2796 : : {Py_tp_iter, PyObject_SelfIter},
2797 : : {Py_tp_iternext, combinations_next},
2798 : : {Py_tp_methods, combinations_methods},
2799 : : {Py_tp_new, itertools_combinations},
2800 : : {Py_tp_free, PyObject_GC_Del},
2801 : : {0, NULL},
2802 : : };
2803 : :
2804 : : static PyType_Spec combinations_spec = {
2805 : : .name = "itertools.combinations",
2806 : : .basicsize = sizeof(combinationsobject),
2807 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
2808 : : Py_TPFLAGS_IMMUTABLETYPE),
2809 : : .slots = combinations_slots,
2810 : : };
2811 : :
2812 : :
2813 : : /* combinations with replacement object **************************************/
2814 : :
2815 : : /* Equivalent to:
2816 : :
2817 : : def combinations_with_replacement(iterable, r):
2818 : : "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
2819 : : # number items returned: (n+r-1)! / r! / (n-1)!
2820 : : pool = tuple(iterable)
2821 : : n = len(pool)
2822 : : indices = [0] * r
2823 : : yield tuple(pool[i] for i in indices)
2824 : : while 1:
2825 : : for i in reversed(range(r)):
2826 : : if indices[i] != n - 1:
2827 : : break
2828 : : else:
2829 : : return
2830 : : indices[i:] = [indices[i] + 1] * (r - i)
2831 : : yield tuple(pool[i] for i in indices)
2832 : :
2833 : : def combinations_with_replacement2(iterable, r):
2834 : : 'Alternate version that filters from product()'
2835 : : pool = tuple(iterable)
2836 : : n = len(pool)
2837 : : for indices in product(range(n), repeat=r):
2838 : : if sorted(indices) == list(indices):
2839 : : yield tuple(pool[i] for i in indices)
2840 : : */
2841 : : typedef struct {
2842 : : PyObject_HEAD
2843 : : PyObject *pool; /* input converted to a tuple */
2844 : : Py_ssize_t *indices; /* one index per result element */
2845 : : PyObject *result; /* most recently returned result tuple */
2846 : : Py_ssize_t r; /* size of result tuple */
2847 : : int stopped; /* set to 1 when the cwr iterator is exhausted */
2848 : : } cwrobject;
2849 : :
2850 : : /*[clinic input]
2851 : : @classmethod
2852 : : itertools.combinations_with_replacement.__new__
2853 : : iterable: object
2854 : : r: Py_ssize_t
2855 : : Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats.
2856 : :
2857 : : combinations_with_replacement('ABC', 2) --> ('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')
2858 : : [clinic start generated code]*/
2859 : :
2860 : : static PyObject *
2861 : 0 : itertools_combinations_with_replacement_impl(PyTypeObject *type,
2862 : : PyObject *iterable,
2863 : : Py_ssize_t r)
2864 : : /*[clinic end generated code: output=48b26856d4e659ca input=1dc58e82a0878fdc]*/
2865 : : {
2866 : : cwrobject *co;
2867 : : Py_ssize_t n;
2868 : 0 : PyObject *pool = NULL;
2869 : 0 : Py_ssize_t *indices = NULL;
2870 : : Py_ssize_t i;
2871 : :
2872 : 0 : pool = PySequence_Tuple(iterable);
2873 [ # # ]: 0 : if (pool == NULL)
2874 : 0 : goto error;
2875 : 0 : n = PyTuple_GET_SIZE(pool);
2876 [ # # ]: 0 : if (r < 0) {
2877 : 0 : PyErr_SetString(PyExc_ValueError, "r must be non-negative");
2878 : 0 : goto error;
2879 : : }
2880 : :
2881 [ # # ]: 0 : indices = PyMem_New(Py_ssize_t, r);
2882 [ # # ]: 0 : if (indices == NULL) {
2883 : 0 : PyErr_NoMemory();
2884 : 0 : goto error;
2885 : : }
2886 : :
2887 [ # # ]: 0 : for (i=0 ; i<r ; i++)
2888 : 0 : indices[i] = 0;
2889 : :
2890 : : /* create cwrobject structure */
2891 : 0 : co = (cwrobject *)type->tp_alloc(type, 0);
2892 [ # # ]: 0 : if (co == NULL)
2893 : 0 : goto error;
2894 : :
2895 : 0 : co->pool = pool;
2896 : 0 : co->indices = indices;
2897 : 0 : co->result = NULL;
2898 : 0 : co->r = r;
2899 [ # # # # ]: 0 : co->stopped = !n && r;
2900 : :
2901 : 0 : return (PyObject *)co;
2902 : :
2903 : 0 : error:
2904 [ # # ]: 0 : if (indices != NULL)
2905 : 0 : PyMem_Free(indices);
2906 : 0 : Py_XDECREF(pool);
2907 : 0 : return NULL;
2908 : : }
2909 : :
2910 : : static void
2911 : 0 : cwr_dealloc(cwrobject *co)
2912 : : {
2913 : 0 : PyTypeObject *tp = Py_TYPE(co);
2914 : 0 : PyObject_GC_UnTrack(co);
2915 : 0 : Py_XDECREF(co->pool);
2916 : 0 : Py_XDECREF(co->result);
2917 [ # # ]: 0 : if (co->indices != NULL)
2918 : 0 : PyMem_Free(co->indices);
2919 : 0 : tp->tp_free(co);
2920 : 0 : Py_DECREF(tp);
2921 : 0 : }
2922 : :
2923 : : static PyObject *
2924 : 0 : cwr_sizeof(cwrobject *co, void *unused)
2925 : : {
2926 : 0 : size_t res = _PyObject_SIZE(Py_TYPE(co));
2927 : 0 : res += (size_t)co->r * sizeof(Py_ssize_t);
2928 : 0 : return PyLong_FromSize_t(res);
2929 : : }
2930 : :
2931 : : static int
2932 : 0 : cwr_traverse(cwrobject *co, visitproc visit, void *arg)
2933 : : {
2934 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(co));
2935 [ # # # # ]: 0 : Py_VISIT(co->pool);
2936 [ # # # # ]: 0 : Py_VISIT(co->result);
2937 : 0 : return 0;
2938 : : }
2939 : :
2940 : : static PyObject *
2941 : 0 : cwr_next(cwrobject *co)
2942 : : {
2943 : : PyObject *elem;
2944 : : PyObject *oldelem;
2945 : 0 : PyObject *pool = co->pool;
2946 : 0 : Py_ssize_t *indices = co->indices;
2947 : 0 : PyObject *result = co->result;
2948 : 0 : Py_ssize_t n = PyTuple_GET_SIZE(pool);
2949 : 0 : Py_ssize_t r = co->r;
2950 : : Py_ssize_t i, index;
2951 : :
2952 [ # # ]: 0 : if (co->stopped)
2953 : 0 : return NULL;
2954 : :
2955 [ # # ]: 0 : if (result == NULL) {
2956 : : /* On the first pass, initialize result tuple with pool[0] */
2957 : 0 : result = PyTuple_New(r);
2958 [ # # ]: 0 : if (result == NULL)
2959 : 0 : goto empty;
2960 : 0 : co->result = result;
2961 [ # # ]: 0 : if (n > 0) {
2962 : 0 : elem = PyTuple_GET_ITEM(pool, 0);
2963 [ # # ]: 0 : for (i=0; i<r ; i++) {
2964 : : assert(indices[i] == 0);
2965 : 0 : Py_INCREF(elem);
2966 : 0 : PyTuple_SET_ITEM(result, i, elem);
2967 : : }
2968 : : }
2969 : : } else {
2970 : : /* Copy the previous result tuple or re-use it if available */
2971 [ # # ]: 0 : if (Py_REFCNT(result) > 1) {
2972 : 0 : PyObject *old_result = result;
2973 : 0 : result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
2974 [ # # ]: 0 : if (result == NULL)
2975 : 0 : goto empty;
2976 : 0 : co->result = result;
2977 : 0 : Py_DECREF(old_result);
2978 : : }
2979 : : // bpo-42536: The GC may have untracked this result tuple. Since we're
2980 : : // recycling it, make sure it's tracked again:
2981 [ # # ]: 0 : else if (!_PyObject_GC_IS_TRACKED(result)) {
2982 : 0 : _PyObject_GC_TRACK(result);
2983 : : }
2984 : : /* Now, we've got the only copy so we can update it in-place CPython's
2985 : : empty tuple is a singleton and cached in PyTuple's freelist. */
2986 : : assert(r == 0 || Py_REFCNT(result) == 1);
2987 : :
2988 : : /* Scan indices right-to-left until finding one that is not
2989 : : * at its maximum (n-1). */
2990 [ # # # # ]: 0 : for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
2991 : : ;
2992 : :
2993 : : /* If i is negative, then the indices are all at
2994 : : their maximum value and we're done. */
2995 [ # # ]: 0 : if (i < 0)
2996 : 0 : goto empty;
2997 : :
2998 : : /* Increment the current index which we know is not at its
2999 : : maximum. Then set all to the right to the same value. */
3000 : 0 : index = indices[i] + 1;
3001 : : assert(index < n);
3002 : 0 : elem = PyTuple_GET_ITEM(pool, index);
3003 [ # # ]: 0 : for ( ; i<r ; i++) {
3004 : 0 : indices[i] = index;
3005 : 0 : Py_INCREF(elem);
3006 : 0 : oldelem = PyTuple_GET_ITEM(result, i);
3007 : 0 : PyTuple_SET_ITEM(result, i, elem);
3008 : 0 : Py_DECREF(oldelem);
3009 : : }
3010 : : }
3011 : :
3012 : 0 : return Py_NewRef(result);
3013 : :
3014 : 0 : empty:
3015 : 0 : co->stopped = 1;
3016 : 0 : return NULL;
3017 : : }
3018 : :
3019 : : static PyObject *
3020 : 0 : cwr_reduce(cwrobject *lz, PyObject *Py_UNUSED(ignored))
3021 : : {
3022 [ # # ]: 0 : if (lz->result == NULL) {
3023 : 0 : return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
3024 [ # # ]: 0 : } else if (lz->stopped) {
3025 : 0 : return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
3026 : : } else {
3027 : : PyObject *indices;
3028 : : Py_ssize_t i;
3029 : :
3030 : : /* we must pickle the indices and use them for setstate */
3031 : 0 : indices = PyTuple_New(lz->r);
3032 [ # # ]: 0 : if (!indices)
3033 : 0 : return NULL;
3034 [ # # ]: 0 : for (i=0; i<lz->r; i++) {
3035 : 0 : PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
3036 [ # # ]: 0 : if (!index) {
3037 : 0 : Py_DECREF(indices);
3038 : 0 : return NULL;
3039 : : }
3040 : 0 : PyTuple_SET_ITEM(indices, i, index);
3041 : : }
3042 : :
3043 : 0 : return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
3044 : : }
3045 : : }
3046 : :
3047 : : static PyObject *
3048 : 0 : cwr_setstate(cwrobject *lz, PyObject *state)
3049 : : {
3050 : : PyObject *result;
3051 : : Py_ssize_t n, i;
3052 : :
3053 [ # # # # ]: 0 : if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
3054 : : {
3055 : 0 : PyErr_SetString(PyExc_ValueError, "invalid arguments");
3056 : 0 : return NULL;
3057 : : }
3058 : :
3059 : 0 : n = PyTuple_GET_SIZE(lz->pool);
3060 [ # # ]: 0 : for (i=0; i<lz->r; i++) {
3061 : 0 : PyObject* indexObject = PyTuple_GET_ITEM(state, i);
3062 : 0 : Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3063 : :
3064 [ # # # # ]: 0 : if (index < 0 && PyErr_Occurred())
3065 : 0 : return NULL; /* not an integer */
3066 : : /* clamp the index */
3067 [ # # ]: 0 : if (index < 0)
3068 : 0 : index = 0;
3069 [ # # ]: 0 : else if (index > n-1)
3070 : 0 : index = n-1;
3071 : 0 : lz->indices[i] = index;
3072 : : }
3073 : 0 : result = PyTuple_New(lz->r);
3074 [ # # ]: 0 : if (result == NULL)
3075 : 0 : return NULL;
3076 [ # # ]: 0 : for (i=0; i<lz->r; i++) {
3077 : 0 : PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
3078 : 0 : Py_INCREF(element);
3079 : 0 : PyTuple_SET_ITEM(result, i, element);
3080 : : }
3081 : 0 : Py_XSETREF(lz->result, result);
3082 : 0 : Py_RETURN_NONE;
3083 : : }
3084 : :
3085 : : static PyMethodDef cwr_methods[] = {
3086 : : {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
3087 : : reduce_doc},
3088 : : {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
3089 : : setstate_doc},
3090 : : {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
3091 : : sizeof_doc},
3092 : : {NULL, NULL} /* sentinel */
3093 : : };
3094 : :
3095 : : static PyType_Slot cwr_slots[] = {
3096 : : {Py_tp_dealloc, cwr_dealloc},
3097 : : {Py_tp_getattro, PyObject_GenericGetAttr},
3098 : : {Py_tp_doc, (void *)itertools_combinations_with_replacement__doc__},
3099 : : {Py_tp_traverse, cwr_traverse},
3100 : : {Py_tp_iter, PyObject_SelfIter},
3101 : : {Py_tp_iternext, cwr_next},
3102 : : {Py_tp_methods, cwr_methods},
3103 : : {Py_tp_new, itertools_combinations_with_replacement},
3104 : : {Py_tp_free, PyObject_GC_Del},
3105 : : {0, NULL},
3106 : : };
3107 : :
3108 : : static PyType_Spec cwr_spec = {
3109 : : .name = "itertools.combinations_with_replacement",
3110 : : .basicsize = sizeof(cwrobject),
3111 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3112 : : Py_TPFLAGS_IMMUTABLETYPE),
3113 : : .slots = cwr_slots,
3114 : : };
3115 : :
3116 : :
3117 : : /* permutations object ********************************************************
3118 : :
3119 : : def permutations(iterable, r=None):
3120 : : # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
3121 : : # permutations(range(3)) --> 012 021 102 120 201 210
3122 : : pool = tuple(iterable)
3123 : : n = len(pool)
3124 : : r = n if r is None else r
3125 : : if r > n:
3126 : : return
3127 : : indices = list(range(n))
3128 : : cycles = list(range(n, n-r, -1))
3129 : : yield tuple(pool[i] for i in indices[:r])
3130 : : while n:
3131 : : for i in reversed(range(r)):
3132 : : cycles[i] -= 1
3133 : : if cycles[i] == 0:
3134 : : indices[i:] = indices[i+1:] + indices[i:i+1]
3135 : : cycles[i] = n - i
3136 : : else:
3137 : : j = cycles[i]
3138 : : indices[i], indices[-j] = indices[-j], indices[i]
3139 : : yield tuple(pool[i] for i in indices[:r])
3140 : : break
3141 : : else:
3142 : : return
3143 : : */
3144 : :
3145 : : typedef struct {
3146 : : PyObject_HEAD
3147 : : PyObject *pool; /* input converted to a tuple */
3148 : : Py_ssize_t *indices; /* one index per element in the pool */
3149 : : Py_ssize_t *cycles; /* one rollover counter per element in the result */
3150 : : PyObject *result; /* most recently returned result tuple */
3151 : : Py_ssize_t r; /* size of result tuple */
3152 : : int stopped; /* set to 1 when the iterator is exhausted */
3153 : : } permutationsobject;
3154 : :
3155 : : /*[clinic input]
3156 : : @classmethod
3157 : : itertools.permutations.__new__
3158 : : iterable: object
3159 : : r as robj: object = None
3160 : : Return successive r-length permutations of elements in the iterable.
3161 : :
3162 : : permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)
3163 : : [clinic start generated code]*/
3164 : :
3165 : : static PyObject *
3166 : 54 : itertools_permutations_impl(PyTypeObject *type, PyObject *iterable,
3167 : : PyObject *robj)
3168 : : /*[clinic end generated code: output=296a72fa76d620ea input=57d0170a4ac0ec7a]*/
3169 : : {
3170 : : permutationsobject *po;
3171 : : Py_ssize_t n;
3172 : : Py_ssize_t r;
3173 : 54 : PyObject *pool = NULL;
3174 : 54 : Py_ssize_t *indices = NULL;
3175 : 54 : Py_ssize_t *cycles = NULL;
3176 : : Py_ssize_t i;
3177 : :
3178 : 54 : pool = PySequence_Tuple(iterable);
3179 [ - + ]: 54 : if (pool == NULL)
3180 : 0 : goto error;
3181 : 54 : n = PyTuple_GET_SIZE(pool);
3182 : :
3183 : 54 : r = n;
3184 [ - + ]: 54 : if (robj != Py_None) {
3185 [ # # ]: 0 : if (!PyLong_Check(robj)) {
3186 : 0 : PyErr_SetString(PyExc_TypeError, "Expected int as r");
3187 : 0 : goto error;
3188 : : }
3189 : 0 : r = PyLong_AsSsize_t(robj);
3190 [ # # # # ]: 0 : if (r == -1 && PyErr_Occurred())
3191 : 0 : goto error;
3192 : : }
3193 [ - + ]: 54 : if (r < 0) {
3194 : 0 : PyErr_SetString(PyExc_ValueError, "r must be non-negative");
3195 : 0 : goto error;
3196 : : }
3197 : :
3198 [ + - ]: 54 : indices = PyMem_New(Py_ssize_t, n);
3199 [ + - ]: 54 : cycles = PyMem_New(Py_ssize_t, r);
3200 [ + - - + ]: 54 : if (indices == NULL || cycles == NULL) {
3201 : 0 : PyErr_NoMemory();
3202 : 0 : goto error;
3203 : : }
3204 : :
3205 [ + + ]: 126 : for (i=0 ; i<n ; i++)
3206 : 72 : indices[i] = i;
3207 [ + + ]: 126 : for (i=0 ; i<r ; i++)
3208 : 72 : cycles[i] = n - i;
3209 : :
3210 : : /* create permutationsobject structure */
3211 : 54 : po = (permutationsobject *)type->tp_alloc(type, 0);
3212 [ - + ]: 54 : if (po == NULL)
3213 : 0 : goto error;
3214 : :
3215 : 54 : po->pool = pool;
3216 : 54 : po->indices = indices;
3217 : 54 : po->cycles = cycles;
3218 : 54 : po->result = NULL;
3219 : 54 : po->r = r;
3220 : 54 : po->stopped = r > n ? 1 : 0;
3221 : :
3222 : 54 : return (PyObject *)po;
3223 : :
3224 : 0 : error:
3225 [ # # ]: 0 : if (indices != NULL)
3226 : 0 : PyMem_Free(indices);
3227 [ # # ]: 0 : if (cycles != NULL)
3228 : 0 : PyMem_Free(cycles);
3229 : 0 : Py_XDECREF(pool);
3230 : 0 : return NULL;
3231 : : }
3232 : :
3233 : : static void
3234 : 54 : permutations_dealloc(permutationsobject *po)
3235 : : {
3236 : 54 : PyTypeObject *tp = Py_TYPE(po);
3237 : 54 : PyObject_GC_UnTrack(po);
3238 : 54 : Py_XDECREF(po->pool);
3239 : 54 : Py_XDECREF(po->result);
3240 : 54 : PyMem_Free(po->indices);
3241 : 54 : PyMem_Free(po->cycles);
3242 : 54 : tp->tp_free(po);
3243 : 54 : Py_DECREF(tp);
3244 : 54 : }
3245 : :
3246 : : static PyObject *
3247 : 0 : permutations_sizeof(permutationsobject *po, void *unused)
3248 : : {
3249 : 0 : size_t res = _PyObject_SIZE(Py_TYPE(po));
3250 : 0 : res += (size_t)PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
3251 : 0 : res += (size_t)po->r * sizeof(Py_ssize_t);
3252 : 0 : return PyLong_FromSize_t(res);
3253 : : }
3254 : :
3255 : : static int
3256 : 0 : permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
3257 : : {
3258 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(po));
3259 [ # # # # ]: 0 : Py_VISIT(po->pool);
3260 [ # # # # ]: 0 : Py_VISIT(po->result);
3261 : 0 : return 0;
3262 : : }
3263 : :
3264 : : static PyObject *
3265 : 126 : permutations_next(permutationsobject *po)
3266 : : {
3267 : : PyObject *elem;
3268 : : PyObject *oldelem;
3269 : 126 : PyObject *pool = po->pool;
3270 : 126 : Py_ssize_t *indices = po->indices;
3271 : 126 : Py_ssize_t *cycles = po->cycles;
3272 : 126 : PyObject *result = po->result;
3273 : 126 : Py_ssize_t n = PyTuple_GET_SIZE(pool);
3274 : 126 : Py_ssize_t r = po->r;
3275 : : Py_ssize_t i, j, k, index;
3276 : :
3277 [ - + ]: 126 : if (po->stopped)
3278 : 0 : return NULL;
3279 : :
3280 [ + + ]: 126 : if (result == NULL) {
3281 : : /* On the first pass, initialize result tuple using the indices */
3282 : 54 : result = PyTuple_New(r);
3283 [ - + ]: 54 : if (result == NULL)
3284 : 0 : goto empty;
3285 : 54 : po->result = result;
3286 [ + + ]: 126 : for (i=0; i<r ; i++) {
3287 : 72 : index = indices[i];
3288 : 72 : elem = PyTuple_GET_ITEM(pool, index);
3289 : 72 : Py_INCREF(elem);
3290 : 72 : PyTuple_SET_ITEM(result, i, elem);
3291 : : }
3292 : : } else {
3293 [ - + ]: 72 : if (n == 0)
3294 : 0 : goto empty;
3295 : :
3296 : : /* Copy the previous result tuple or re-use it if available */
3297 [ + - ]: 72 : if (Py_REFCNT(result) > 1) {
3298 : 72 : PyObject *old_result = result;
3299 : 72 : result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
3300 [ - + ]: 72 : if (result == NULL)
3301 : 0 : goto empty;
3302 : 72 : po->result = result;
3303 : 72 : Py_DECREF(old_result);
3304 : : }
3305 : : // bpo-42536: The GC may have untracked this result tuple. Since we're
3306 : : // recycling it, make sure it's tracked again:
3307 [ # # ]: 0 : else if (!_PyObject_GC_IS_TRACKED(result)) {
3308 : 0 : _PyObject_GC_TRACK(result);
3309 : : }
3310 : : /* Now, we've got the only copy so we can update it in-place */
3311 : : assert(r == 0 || Py_REFCNT(result) == 1);
3312 : :
3313 : : /* Decrement rightmost cycle, moving leftward upon zero rollover */
3314 [ + + ]: 162 : for (i=r-1 ; i>=0 ; i--) {
3315 : 108 : cycles[i] -= 1;
3316 [ + + ]: 108 : if (cycles[i] == 0) {
3317 : : /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
3318 : 90 : index = indices[i];
3319 [ + + ]: 108 : for (j=i ; j<n-1 ; j++)
3320 : 18 : indices[j] = indices[j+1];
3321 : 90 : indices[n-1] = index;
3322 : 90 : cycles[i] = n - i;
3323 : : } else {
3324 : 18 : j = cycles[i];
3325 : 18 : index = indices[i];
3326 : 18 : indices[i] = indices[n-j];
3327 : 18 : indices[n-j] = index;
3328 : :
3329 [ + + ]: 54 : for (k=i; k<r ; k++) {
3330 : : /* start with i, the leftmost element that changed */
3331 : : /* yield tuple(pool[k] for k in indices[:r]) */
3332 : 36 : index = indices[k];
3333 : 36 : elem = PyTuple_GET_ITEM(pool, index);
3334 : 36 : Py_INCREF(elem);
3335 : 36 : oldelem = PyTuple_GET_ITEM(result, k);
3336 : 36 : PyTuple_SET_ITEM(result, k, elem);
3337 : 36 : Py_DECREF(oldelem);
3338 : : }
3339 : 18 : break;
3340 : : }
3341 : : }
3342 : : /* If i is negative, then the cycles have all
3343 : : rolled-over and we're done. */
3344 [ + + ]: 72 : if (i < 0)
3345 : 54 : goto empty;
3346 : : }
3347 : 72 : return Py_NewRef(result);
3348 : :
3349 : 54 : empty:
3350 : 54 : po->stopped = 1;
3351 : 54 : return NULL;
3352 : : }
3353 : :
3354 : : static PyObject *
3355 : 0 : permutations_reduce(permutationsobject *po, PyObject *Py_UNUSED(ignored))
3356 : : {
3357 [ # # ]: 0 : if (po->result == NULL) {
3358 : 0 : return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
3359 [ # # ]: 0 : } else if (po->stopped) {
3360 : 0 : return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
3361 : : } else {
3362 : 0 : PyObject *indices=NULL, *cycles=NULL;
3363 : : Py_ssize_t n, i;
3364 : :
3365 : : /* we must pickle the indices and cycles and use them for setstate */
3366 : 0 : n = PyTuple_GET_SIZE(po->pool);
3367 : 0 : indices = PyTuple_New(n);
3368 [ # # ]: 0 : if (indices == NULL)
3369 : 0 : goto err;
3370 [ # # ]: 0 : for (i=0; i<n; i++) {
3371 : 0 : PyObject* index = PyLong_FromSsize_t(po->indices[i]);
3372 [ # # ]: 0 : if (!index)
3373 : 0 : goto err;
3374 : 0 : PyTuple_SET_ITEM(indices, i, index);
3375 : : }
3376 : :
3377 : 0 : cycles = PyTuple_New(po->r);
3378 [ # # ]: 0 : if (cycles == NULL)
3379 : 0 : goto err;
3380 [ # # ]: 0 : for (i=0 ; i<po->r ; i++) {
3381 : 0 : PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
3382 [ # # ]: 0 : if (!index)
3383 : 0 : goto err;
3384 : 0 : PyTuple_SET_ITEM(cycles, i, index);
3385 : : }
3386 : 0 : return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
3387 : : po->pool, po->r,
3388 : : indices, cycles);
3389 : 0 : err:
3390 : 0 : Py_XDECREF(indices);
3391 : 0 : Py_XDECREF(cycles);
3392 : 0 : return NULL;
3393 : : }
3394 : : }
3395 : :
3396 : : static PyObject *
3397 : 0 : permutations_setstate(permutationsobject *po, PyObject *state)
3398 : : {
3399 : : PyObject *indices, *cycles, *result;
3400 : : Py_ssize_t n, i;
3401 : :
3402 [ # # ]: 0 : if (!PyTuple_Check(state)) {
3403 : 0 : PyErr_SetString(PyExc_TypeError, "state is not a tuple");
3404 : 0 : return NULL;
3405 : : }
3406 [ # # ]: 0 : if (!PyArg_ParseTuple(state, "O!O!",
3407 : : &PyTuple_Type, &indices,
3408 : : &PyTuple_Type, &cycles)) {
3409 : 0 : return NULL;
3410 : : }
3411 : :
3412 : 0 : n = PyTuple_GET_SIZE(po->pool);
3413 [ # # # # ]: 0 : if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
3414 : 0 : PyErr_SetString(PyExc_ValueError, "invalid arguments");
3415 : 0 : return NULL;
3416 : : }
3417 : :
3418 [ # # ]: 0 : for (i=0; i<n; i++) {
3419 : 0 : PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
3420 : 0 : Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3421 [ # # # # ]: 0 : if (index < 0 && PyErr_Occurred())
3422 : 0 : return NULL; /* not an integer */
3423 : : /* clamp the index */
3424 [ # # ]: 0 : if (index < 0)
3425 : 0 : index = 0;
3426 [ # # ]: 0 : else if (index > n-1)
3427 : 0 : index = n-1;
3428 : 0 : po->indices[i] = index;
3429 : : }
3430 : :
3431 [ # # ]: 0 : for (i=0; i<po->r; i++) {
3432 : 0 : PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
3433 : 0 : Py_ssize_t index = PyLong_AsSsize_t(indexObject);
3434 [ # # # # ]: 0 : if (index < 0 && PyErr_Occurred())
3435 : 0 : return NULL; /* not an integer */
3436 [ # # ]: 0 : if (index < 1)
3437 : 0 : index = 1;
3438 [ # # ]: 0 : else if (index > n-i)
3439 : 0 : index = n-i;
3440 : 0 : po->cycles[i] = index;
3441 : : }
3442 : 0 : result = PyTuple_New(po->r);
3443 [ # # ]: 0 : if (result == NULL)
3444 : 0 : return NULL;
3445 [ # # ]: 0 : for (i=0; i<po->r; i++) {
3446 : 0 : PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
3447 : 0 : Py_INCREF(element);
3448 : 0 : PyTuple_SET_ITEM(result, i, element);
3449 : : }
3450 : 0 : Py_XSETREF(po->result, result);
3451 : 0 : Py_RETURN_NONE;
3452 : : }
3453 : :
3454 : : static PyMethodDef permuations_methods[] = {
3455 : : {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
3456 : : reduce_doc},
3457 : : {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
3458 : : setstate_doc},
3459 : : {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
3460 : : sizeof_doc},
3461 : : {NULL, NULL} /* sentinel */
3462 : : };
3463 : :
3464 : : static PyType_Slot permutations_slots[] = {
3465 : : {Py_tp_dealloc, permutations_dealloc},
3466 : : {Py_tp_getattro, PyObject_GenericGetAttr},
3467 : : {Py_tp_doc, (void *)itertools_permutations__doc__},
3468 : : {Py_tp_traverse, permutations_traverse},
3469 : : {Py_tp_iter, PyObject_SelfIter},
3470 : : {Py_tp_iternext, permutations_next},
3471 : : {Py_tp_methods, permuations_methods},
3472 : : {Py_tp_new, itertools_permutations},
3473 : : {Py_tp_free, PyObject_GC_Del},
3474 : : {0, NULL},
3475 : : };
3476 : :
3477 : : static PyType_Spec permutations_spec = {
3478 : : .name = "itertools.permutations",
3479 : : .basicsize = sizeof(permutationsobject),
3480 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3481 : : Py_TPFLAGS_IMMUTABLETYPE),
3482 : : .slots = permutations_slots,
3483 : : };
3484 : :
3485 : :
3486 : : /* accumulate object ********************************************************/
3487 : :
3488 : : typedef struct {
3489 : : PyObject_HEAD
3490 : : PyObject *total;
3491 : : PyObject *it;
3492 : : PyObject *binop;
3493 : : PyObject *initial;
3494 : : itertools_state *state;
3495 : : } accumulateobject;
3496 : :
3497 : : /*[clinic input]
3498 : : @classmethod
3499 : : itertools.accumulate.__new__
3500 : : iterable: object
3501 : : func as binop: object = None
3502 : : *
3503 : : initial: object = None
3504 : : Return series of accumulated sums (or other binary function results).
3505 : : [clinic start generated code]*/
3506 : :
3507 : : static PyObject *
3508 : 0 : itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
3509 : : PyObject *binop, PyObject *initial)
3510 : : /*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
3511 : : {
3512 : : PyObject *it;
3513 : : accumulateobject *lz;
3514 : :
3515 : : /* Get iterator. */
3516 : 0 : it = PyObject_GetIter(iterable);
3517 [ # # ]: 0 : if (it == NULL)
3518 : 0 : return NULL;
3519 : :
3520 : : /* create accumulateobject structure */
3521 : 0 : lz = (accumulateobject *)type->tp_alloc(type, 0);
3522 [ # # ]: 0 : if (lz == NULL) {
3523 : 0 : Py_DECREF(it);
3524 : 0 : return NULL;
3525 : : }
3526 : :
3527 [ # # ]: 0 : if (binop != Py_None) {
3528 : 0 : lz->binop = Py_XNewRef(binop);
3529 : : }
3530 : 0 : lz->total = NULL;
3531 : 0 : lz->it = it;
3532 : 0 : lz->initial = Py_XNewRef(initial);
3533 : 0 : lz->state = find_state_by_type(type);
3534 : 0 : return (PyObject *)lz;
3535 : : }
3536 : :
3537 : : static void
3538 : 0 : accumulate_dealloc(accumulateobject *lz)
3539 : : {
3540 : 0 : PyTypeObject *tp = Py_TYPE(lz);
3541 : 0 : PyObject_GC_UnTrack(lz);
3542 : 0 : Py_XDECREF(lz->binop);
3543 : 0 : Py_XDECREF(lz->total);
3544 : 0 : Py_XDECREF(lz->it);
3545 : 0 : Py_XDECREF(lz->initial);
3546 : 0 : tp->tp_free(lz);
3547 : 0 : Py_DECREF(tp);
3548 : 0 : }
3549 : :
3550 : : static int
3551 : 0 : accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
3552 : : {
3553 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(lz));
3554 [ # # # # ]: 0 : Py_VISIT(lz->binop);
3555 [ # # # # ]: 0 : Py_VISIT(lz->it);
3556 [ # # # # ]: 0 : Py_VISIT(lz->total);
3557 [ # # # # ]: 0 : Py_VISIT(lz->initial);
3558 : 0 : return 0;
3559 : : }
3560 : :
3561 : : static PyObject *
3562 : 0 : accumulate_next(accumulateobject *lz)
3563 : : {
3564 : : PyObject *val, *newtotal;
3565 : :
3566 [ # # ]: 0 : if (lz->initial != Py_None) {
3567 : 0 : lz->total = lz->initial;
3568 : 0 : lz->initial = Py_NewRef(Py_None);
3569 : 0 : return Py_NewRef(lz->total);
3570 : : }
3571 : 0 : val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
3572 [ # # ]: 0 : if (val == NULL)
3573 : 0 : return NULL;
3574 : :
3575 [ # # ]: 0 : if (lz->total == NULL) {
3576 : 0 : lz->total = Py_NewRef(val);
3577 : 0 : return lz->total;
3578 : : }
3579 : :
3580 [ # # ]: 0 : if (lz->binop == NULL)
3581 : 0 : newtotal = PyNumber_Add(lz->total, val);
3582 : : else
3583 : 0 : newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
3584 : 0 : Py_DECREF(val);
3585 [ # # ]: 0 : if (newtotal == NULL)
3586 : 0 : return NULL;
3587 : :
3588 : 0 : Py_INCREF(newtotal);
3589 : 0 : Py_SETREF(lz->total, newtotal);
3590 : 0 : return newtotal;
3591 : : }
3592 : :
3593 : : static PyObject *
3594 : 0 : accumulate_reduce(accumulateobject *lz, PyObject *Py_UNUSED(ignored))
3595 : : {
3596 : 0 : itertools_state *state = lz->state;
3597 : :
3598 [ # # ]: 0 : if (lz->initial != Py_None) {
3599 : : PyObject *it;
3600 : :
3601 : : assert(lz->total == NULL);
3602 : 0 : it = PyObject_CallFunction((PyObject *)(state->chain_type), "(O)O",
3603 : : lz->initial, lz->it);
3604 [ # # ]: 0 : if (it == NULL)
3605 : 0 : return NULL;
3606 : 0 : return Py_BuildValue("O(NO)O", Py_TYPE(lz),
3607 [ # # ]: 0 : it, lz->binop?lz->binop:Py_None, Py_None);
3608 : : }
3609 [ # # ]: 0 : if (lz->total == Py_None) {
3610 : : PyObject *it;
3611 : :
3612 : 0 : it = PyObject_CallFunction((PyObject *)(state->chain_type), "(O)O",
3613 : : lz->total, lz->it);
3614 [ # # ]: 0 : if (it == NULL)
3615 : 0 : return NULL;
3616 : 0 : it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
3617 [ # # ]: 0 : it, lz->binop ? lz->binop : Py_None);
3618 [ # # ]: 0 : if (it == NULL)
3619 : 0 : return NULL;
3620 : :
3621 : 0 : return Py_BuildValue("O(NiO)", state->islice_type, it, 1, Py_None);
3622 : : }
3623 : 0 : return Py_BuildValue("O(OO)O", Py_TYPE(lz),
3624 [ # # ]: 0 : lz->it, lz->binop?lz->binop:Py_None,
3625 [ # # ]: 0 : lz->total?lz->total:Py_None);
3626 : : }
3627 : :
3628 : : static PyObject *
3629 : 0 : accumulate_setstate(accumulateobject *lz, PyObject *state)
3630 : : {
3631 : 0 : Py_INCREF(state);
3632 : 0 : Py_XSETREF(lz->total, state);
3633 : 0 : Py_RETURN_NONE;
3634 : : }
3635 : :
3636 : : static PyMethodDef accumulate_methods[] = {
3637 : : {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
3638 : : reduce_doc},
3639 : : {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
3640 : : setstate_doc},
3641 : : {NULL, NULL} /* sentinel */
3642 : : };
3643 : :
3644 : : static PyType_Slot accumulate_slots[] = {
3645 : : {Py_tp_dealloc, accumulate_dealloc},
3646 : : {Py_tp_getattro, PyObject_GenericGetAttr},
3647 : : {Py_tp_doc, (void *)itertools_accumulate__doc__},
3648 : : {Py_tp_traverse, accumulate_traverse},
3649 : : {Py_tp_iter, PyObject_SelfIter},
3650 : : {Py_tp_iternext, accumulate_next},
3651 : : {Py_tp_methods, accumulate_methods},
3652 : : {Py_tp_new, itertools_accumulate},
3653 : : {Py_tp_free, PyObject_GC_Del},
3654 : : {0, NULL},
3655 : : };
3656 : :
3657 : : static PyType_Spec accumulate_spec = {
3658 : : .name = "itertools.accumulate",
3659 : : .basicsize = sizeof(accumulateobject),
3660 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3661 : : Py_TPFLAGS_IMMUTABLETYPE),
3662 : : .slots = accumulate_slots,
3663 : : };
3664 : :
3665 : :
3666 : : /* compress object ************************************************************/
3667 : :
3668 : : /* Equivalent to:
3669 : :
3670 : : def compress(data, selectors):
3671 : : "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
3672 : : return (d for d, s in zip(data, selectors) if s)
3673 : : */
3674 : :
3675 : : typedef struct {
3676 : : PyObject_HEAD
3677 : : PyObject *data;
3678 : : PyObject *selectors;
3679 : : } compressobject;
3680 : :
3681 : : /*[clinic input]
3682 : : @classmethod
3683 : : itertools.compress.__new__
3684 : : data as seq1: object
3685 : : selectors as seq2: object
3686 : : Return data elements corresponding to true selector elements.
3687 : :
3688 : : Forms a shorter iterator from selected data elements using the selectors to
3689 : : choose the data elements.
3690 : : [clinic start generated code]*/
3691 : :
3692 : : static PyObject *
3693 : 0 : itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
3694 : : /*[clinic end generated code: output=7e67157212ed09e0 input=79596d7cd20c77e5]*/
3695 : : {
3696 : 0 : PyObject *data=NULL, *selectors=NULL;
3697 : : compressobject *lz;
3698 : :
3699 : 0 : data = PyObject_GetIter(seq1);
3700 [ # # ]: 0 : if (data == NULL)
3701 : 0 : goto fail;
3702 : 0 : selectors = PyObject_GetIter(seq2);
3703 [ # # ]: 0 : if (selectors == NULL)
3704 : 0 : goto fail;
3705 : :
3706 : : /* create compressobject structure */
3707 : 0 : lz = (compressobject *)type->tp_alloc(type, 0);
3708 [ # # ]: 0 : if (lz == NULL)
3709 : 0 : goto fail;
3710 : 0 : lz->data = data;
3711 : 0 : lz->selectors = selectors;
3712 : 0 : return (PyObject *)lz;
3713 : :
3714 : 0 : fail:
3715 : 0 : Py_XDECREF(data);
3716 : 0 : Py_XDECREF(selectors);
3717 : 0 : return NULL;
3718 : : }
3719 : :
3720 : : static void
3721 : 0 : compress_dealloc(compressobject *lz)
3722 : : {
3723 : 0 : PyTypeObject *tp = Py_TYPE(lz);
3724 : 0 : PyObject_GC_UnTrack(lz);
3725 : 0 : Py_XDECREF(lz->data);
3726 : 0 : Py_XDECREF(lz->selectors);
3727 : 0 : tp->tp_free(lz);
3728 : 0 : Py_DECREF(tp);
3729 : 0 : }
3730 : :
3731 : : static int
3732 : 0 : compress_traverse(compressobject *lz, visitproc visit, void *arg)
3733 : : {
3734 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(lz));
3735 [ # # # # ]: 0 : Py_VISIT(lz->data);
3736 [ # # # # ]: 0 : Py_VISIT(lz->selectors);
3737 : 0 : return 0;
3738 : : }
3739 : :
3740 : : static PyObject *
3741 : 0 : compress_next(compressobject *lz)
3742 : : {
3743 : 0 : PyObject *data = lz->data, *selectors = lz->selectors;
3744 : : PyObject *datum, *selector;
3745 : 0 : PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
3746 : 0 : PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
3747 : : int ok;
3748 : :
3749 : : while (1) {
3750 : : /* Steps: get datum, get selector, evaluate selector.
3751 : : Order is important (to match the pure python version
3752 : : in terms of which input gets a chance to raise an
3753 : : exception first).
3754 : : */
3755 : :
3756 : 0 : datum = datanext(data);
3757 [ # # ]: 0 : if (datum == NULL)
3758 : 0 : return NULL;
3759 : :
3760 : 0 : selector = selectornext(selectors);
3761 [ # # ]: 0 : if (selector == NULL) {
3762 : 0 : Py_DECREF(datum);
3763 : 0 : return NULL;
3764 : : }
3765 : :
3766 : 0 : ok = PyObject_IsTrue(selector);
3767 : 0 : Py_DECREF(selector);
3768 [ # # ]: 0 : if (ok > 0)
3769 : 0 : return datum;
3770 : 0 : Py_DECREF(datum);
3771 [ # # ]: 0 : if (ok < 0)
3772 : 0 : return NULL;
3773 : : }
3774 : : }
3775 : :
3776 : : static PyObject *
3777 : 0 : compress_reduce(compressobject *lz, PyObject *Py_UNUSED(ignored))
3778 : : {
3779 : 0 : return Py_BuildValue("O(OO)", Py_TYPE(lz),
3780 : : lz->data, lz->selectors);
3781 : : }
3782 : :
3783 : : static PyMethodDef compress_methods[] = {
3784 : : {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
3785 : : reduce_doc},
3786 : : {NULL, NULL} /* sentinel */
3787 : : };
3788 : :
3789 : : static PyType_Slot compress_slots[] = {
3790 : : {Py_tp_dealloc, compress_dealloc},
3791 : : {Py_tp_getattro, PyObject_GenericGetAttr},
3792 : : {Py_tp_doc, (void *)itertools_compress__doc__},
3793 : : {Py_tp_traverse, compress_traverse},
3794 : : {Py_tp_iter, PyObject_SelfIter},
3795 : : {Py_tp_iternext, compress_next},
3796 : : {Py_tp_methods, compress_methods},
3797 : : {Py_tp_new, itertools_compress},
3798 : : {Py_tp_free, PyObject_GC_Del},
3799 : : {0, NULL},
3800 : : };
3801 : :
3802 : : static PyType_Spec compress_spec = {
3803 : : .name = "itertools.compress",
3804 : : .basicsize = sizeof(compressobject),
3805 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3806 : : Py_TPFLAGS_IMMUTABLETYPE),
3807 : : .slots = compress_slots,
3808 : : };
3809 : :
3810 : :
3811 : : /* filterfalse object ************************************************************/
3812 : :
3813 : : typedef struct {
3814 : : PyObject_HEAD
3815 : : PyObject *func;
3816 : : PyObject *it;
3817 : : } filterfalseobject;
3818 : :
3819 : : /*[clinic input]
3820 : : @classmethod
3821 : : itertools.filterfalse.__new__
3822 : : function as func: object
3823 : : iterable as seq: object
3824 : : /
3825 : : Return those items of iterable for which function(item) is false.
3826 : :
3827 : : If function is None, return the items that are false.
3828 : : [clinic start generated code]*/
3829 : :
3830 : : static PyObject *
3831 : 0 : itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
3832 : : /*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
3833 : : {
3834 : : PyObject *it;
3835 : : filterfalseobject *lz;
3836 : :
3837 : : /* Get iterator. */
3838 : 0 : it = PyObject_GetIter(seq);
3839 [ # # ]: 0 : if (it == NULL)
3840 : 0 : return NULL;
3841 : :
3842 : : /* create filterfalseobject structure */
3843 : 0 : lz = (filterfalseobject *)type->tp_alloc(type, 0);
3844 [ # # ]: 0 : if (lz == NULL) {
3845 : 0 : Py_DECREF(it);
3846 : 0 : return NULL;
3847 : : }
3848 : 0 : lz->func = Py_NewRef(func);
3849 : 0 : lz->it = it;
3850 : :
3851 : 0 : return (PyObject *)lz;
3852 : : }
3853 : :
3854 : : static void
3855 : 0 : filterfalse_dealloc(filterfalseobject *lz)
3856 : : {
3857 : 0 : PyTypeObject *tp = Py_TYPE(lz);
3858 : 0 : PyObject_GC_UnTrack(lz);
3859 : 0 : Py_XDECREF(lz->func);
3860 : 0 : Py_XDECREF(lz->it);
3861 : 0 : tp->tp_free(lz);
3862 : 0 : Py_DECREF(tp);
3863 : 0 : }
3864 : :
3865 : : static int
3866 : 0 : filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
3867 : : {
3868 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(lz));
3869 [ # # # # ]: 0 : Py_VISIT(lz->it);
3870 [ # # # # ]: 0 : Py_VISIT(lz->func);
3871 : 0 : return 0;
3872 : : }
3873 : :
3874 : : static PyObject *
3875 : 0 : filterfalse_next(filterfalseobject *lz)
3876 : : {
3877 : : PyObject *item;
3878 : 0 : PyObject *it = lz->it;
3879 : : long ok;
3880 : : PyObject *(*iternext)(PyObject *);
3881 : :
3882 : 0 : iternext = *Py_TYPE(it)->tp_iternext;
3883 : : for (;;) {
3884 : 0 : item = iternext(it);
3885 [ # # ]: 0 : if (item == NULL)
3886 : 0 : return NULL;
3887 : :
3888 [ # # # # ]: 0 : if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
3889 : 0 : ok = PyObject_IsTrue(item);
3890 : : } else {
3891 : : PyObject *good;
3892 : 0 : good = PyObject_CallOneArg(lz->func, item);
3893 [ # # ]: 0 : if (good == NULL) {
3894 : 0 : Py_DECREF(item);
3895 : 0 : return NULL;
3896 : : }
3897 : 0 : ok = PyObject_IsTrue(good);
3898 : 0 : Py_DECREF(good);
3899 : : }
3900 [ # # ]: 0 : if (ok == 0)
3901 : 0 : return item;
3902 : 0 : Py_DECREF(item);
3903 [ # # ]: 0 : if (ok < 0)
3904 : 0 : return NULL;
3905 : : }
3906 : : }
3907 : :
3908 : : static PyObject *
3909 : 0 : filterfalse_reduce(filterfalseobject *lz, PyObject *Py_UNUSED(ignored))
3910 : : {
3911 : 0 : return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
3912 : : }
3913 : :
3914 : : static PyMethodDef filterfalse_methods[] = {
3915 : : {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
3916 : : reduce_doc},
3917 : : {NULL, NULL} /* sentinel */
3918 : : };
3919 : :
3920 : : static PyType_Slot filterfalse_slots[] = {
3921 : : {Py_tp_dealloc, filterfalse_dealloc},
3922 : : {Py_tp_getattro, PyObject_GenericGetAttr},
3923 : : {Py_tp_doc, (void *)itertools_filterfalse__doc__},
3924 : : {Py_tp_traverse, filterfalse_traverse},
3925 : : {Py_tp_iter, PyObject_SelfIter},
3926 : : {Py_tp_iternext, filterfalse_next},
3927 : : {Py_tp_methods, filterfalse_methods},
3928 : : {Py_tp_new, itertools_filterfalse},
3929 : : {Py_tp_free, PyObject_GC_Del},
3930 : : {0, NULL},
3931 : : };
3932 : :
3933 : : static PyType_Spec filterfalse_spec = {
3934 : : .name = "itertools.filterfalse",
3935 : : .basicsize = sizeof(filterfalseobject),
3936 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
3937 : : Py_TPFLAGS_IMMUTABLETYPE),
3938 : : .slots = filterfalse_slots,
3939 : : };
3940 : :
3941 : :
3942 : : /* count object ************************************************************/
3943 : :
3944 : : typedef struct {
3945 : : PyObject_HEAD
3946 : : Py_ssize_t cnt;
3947 : : PyObject *long_cnt;
3948 : : PyObject *long_step;
3949 : : } countobject;
3950 : :
3951 : : /* Counting logic and invariants:
3952 : :
3953 : : fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
3954 : :
3955 : : assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
3956 : : Advances with: cnt += 1
3957 : : When count hits Y_SSIZE_T_MAX, switch to slow_mode.
3958 : :
3959 : : slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
3960 : :
3961 : : assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
3962 : : All counting is done with python objects (no overflows or underflows).
3963 : : Advances with: long_cnt += long_step
3964 : : Step may be zero -- effectively a slow version of repeat(cnt).
3965 : : Either long_cnt or long_step may be a float, Fraction, or Decimal.
3966 : : */
3967 : :
3968 : : /*[clinic input]
3969 : : @classmethod
3970 : : itertools.count.__new__
3971 : : start as long_cnt: object(c_default="NULL") = 0
3972 : : step as long_step: object(c_default="NULL") = 1
3973 : : Return a count object whose .__next__() method returns consecutive values.
3974 : :
3975 : : Equivalent to:
3976 : : def count(firstval=0, step=1):
3977 : : x = firstval
3978 : : while 1:
3979 : : yield x
3980 : : x += step
3981 : : [clinic start generated code]*/
3982 : :
3983 : : static PyObject *
3984 : 8 : itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
3985 : : PyObject *long_step)
3986 : : /*[clinic end generated code: output=09a9250aebd00b1c input=d7a85eec18bfcd94]*/
3987 : : {
3988 : : countobject *lz;
3989 : : int fast_mode;
3990 : 8 : Py_ssize_t cnt = 0;
3991 : : long step;
3992 : :
3993 [ + + + - : 8 : if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
- + ]
3994 [ # # ]: 0 : (long_step != NULL && !PyNumber_Check(long_step))) {
3995 : 0 : PyErr_SetString(PyExc_TypeError, "a number is required");
3996 : 0 : return NULL;
3997 : : }
3998 : :
3999 [ + + + - : 8 : fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
- + ]
4000 [ # # ]: 0 : (long_step == NULL || PyLong_Check(long_step));
4001 : :
4002 : : /* If not specified, start defaults to 0 */
4003 [ + + ]: 8 : if (long_cnt != NULL) {
4004 [ + - ]: 5 : if (fast_mode) {
4005 : : assert(PyLong_Check(long_cnt));
4006 : 5 : cnt = PyLong_AsSsize_t(long_cnt);
4007 [ - + - - ]: 5 : if (cnt == -1 && PyErr_Occurred()) {
4008 : 0 : PyErr_Clear();
4009 : 0 : fast_mode = 0;
4010 : : }
4011 : : }
4012 : : } else {
4013 : 3 : cnt = 0;
4014 : 3 : long_cnt = _PyLong_GetZero();
4015 : : }
4016 : 8 : Py_INCREF(long_cnt);
4017 : :
4018 : : /* If not specified, step defaults to 1 */
4019 [ + - ]: 8 : if (long_step == NULL) {
4020 : 8 : long_step = _PyLong_GetOne();
4021 : : }
4022 : 8 : Py_INCREF(long_step);
4023 : :
4024 : : assert(long_cnt != NULL && long_step != NULL);
4025 : :
4026 : : /* Fast mode only works when the step is 1 */
4027 [ + - ]: 8 : if (fast_mode) {
4028 : : assert(PyLong_Check(long_step));
4029 : 8 : step = PyLong_AsLong(long_step);
4030 [ - + ]: 8 : if (step != 1) {
4031 : 0 : fast_mode = 0;
4032 [ # # # # ]: 0 : if (step == -1 && PyErr_Occurred())
4033 : 0 : PyErr_Clear();
4034 : : }
4035 : : }
4036 : :
4037 [ + - ]: 8 : if (fast_mode)
4038 [ + - ]: 8 : Py_CLEAR(long_cnt);
4039 : : else
4040 : 0 : cnt = PY_SSIZE_T_MAX;
4041 : :
4042 : : assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
4043 : : (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
4044 : : assert(!fast_mode ||
4045 : : (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
4046 : :
4047 : : /* create countobject structure */
4048 : 8 : lz = (countobject *)type->tp_alloc(type, 0);
4049 [ - + ]: 8 : if (lz == NULL) {
4050 : 0 : Py_XDECREF(long_cnt);
4051 : 0 : Py_DECREF(long_step);
4052 : 0 : return NULL;
4053 : : }
4054 : 8 : lz->cnt = cnt;
4055 : 8 : lz->long_cnt = long_cnt;
4056 : 8 : lz->long_step = long_step;
4057 : :
4058 : 8 : return (PyObject *)lz;
4059 : : }
4060 : :
4061 : : static void
4062 : 8 : count_dealloc(countobject *lz)
4063 : : {
4064 : 8 : PyTypeObject *tp = Py_TYPE(lz);
4065 : 8 : PyObject_GC_UnTrack(lz);
4066 : 8 : Py_XDECREF(lz->long_cnt);
4067 : 8 : Py_XDECREF(lz->long_step);
4068 : 8 : tp->tp_free(lz);
4069 : 8 : Py_DECREF(tp);
4070 : 8 : }
4071 : :
4072 : : static int
4073 : 124 : count_traverse(countobject *lz, visitproc visit, void *arg)
4074 : : {
4075 [ + - - + ]: 124 : Py_VISIT(Py_TYPE(lz));
4076 [ - + - - ]: 124 : Py_VISIT(lz->long_cnt);
4077 [ + - - + ]: 124 : Py_VISIT(lz->long_step);
4078 : 124 : return 0;
4079 : : }
4080 : :
4081 : : static PyObject *
4082 : 0 : count_nextlong(countobject *lz)
4083 : : {
4084 : : PyObject *long_cnt;
4085 : : PyObject *stepped_up;
4086 : :
4087 : 0 : long_cnt = lz->long_cnt;
4088 [ # # ]: 0 : if (long_cnt == NULL) {
4089 : : /* Switch to slow_mode */
4090 : 0 : long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
4091 [ # # ]: 0 : if (long_cnt == NULL)
4092 : 0 : return NULL;
4093 : : }
4094 : : assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
4095 : :
4096 : 0 : stepped_up = PyNumber_Add(long_cnt, lz->long_step);
4097 [ # # ]: 0 : if (stepped_up == NULL)
4098 : 0 : return NULL;
4099 : 0 : lz->long_cnt = stepped_up;
4100 : 0 : return long_cnt;
4101 : : }
4102 : :
4103 : : static PyObject *
4104 : 0 : count_next(countobject *lz)
4105 : : {
4106 [ # # ]: 0 : if (lz->cnt == PY_SSIZE_T_MAX)
4107 : 0 : return count_nextlong(lz);
4108 : 0 : return PyLong_FromSsize_t(lz->cnt++);
4109 : : }
4110 : :
4111 : : static PyObject *
4112 : 0 : count_repr(countobject *lz)
4113 : : {
4114 [ # # ]: 0 : if (lz->cnt != PY_SSIZE_T_MAX)
4115 : 0 : return PyUnicode_FromFormat("%s(%zd)",
4116 : : _PyType_Name(Py_TYPE(lz)), lz->cnt);
4117 : :
4118 [ # # ]: 0 : if (PyLong_Check(lz->long_step)) {
4119 : 0 : long step = PyLong_AsLong(lz->long_step);
4120 [ # # # # ]: 0 : if (step == -1 && PyErr_Occurred()) {
4121 : 0 : PyErr_Clear();
4122 : : }
4123 [ # # ]: 0 : if (step == 1) {
4124 : : /* Don't display step when it is an integer equal to 1 */
4125 : 0 : return PyUnicode_FromFormat("%s(%R)",
4126 : : _PyType_Name(Py_TYPE(lz)),
4127 : : lz->long_cnt);
4128 : : }
4129 : : }
4130 : 0 : return PyUnicode_FromFormat("%s(%R, %R)",
4131 : : _PyType_Name(Py_TYPE(lz)),
4132 : : lz->long_cnt, lz->long_step);
4133 : : }
4134 : :
4135 : : static PyObject *
4136 : 0 : count_reduce(countobject *lz, PyObject *Py_UNUSED(ignored))
4137 : : {
4138 [ # # ]: 0 : if (lz->cnt == PY_SSIZE_T_MAX)
4139 : 0 : return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
4140 : 0 : return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
4141 : : }
4142 : :
4143 : : static PyMethodDef count_methods[] = {
4144 : : {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
4145 : : reduce_doc},
4146 : : {NULL, NULL} /* sentinel */
4147 : : };
4148 : :
4149 : : static PyType_Slot count_slots[] = {
4150 : : {Py_tp_dealloc, count_dealloc},
4151 : : {Py_tp_repr, count_repr},
4152 : : {Py_tp_getattro, PyObject_GenericGetAttr},
4153 : : {Py_tp_doc, (void *)itertools_count__doc__},
4154 : : {Py_tp_traverse, count_traverse},
4155 : : {Py_tp_iter, PyObject_SelfIter},
4156 : : {Py_tp_iternext, count_next},
4157 : : {Py_tp_methods, count_methods},
4158 : : {Py_tp_new, itertools_count},
4159 : : {Py_tp_free, PyObject_GC_Del},
4160 : : {0, NULL},
4161 : : };
4162 : :
4163 : : static PyType_Spec count_spec = {
4164 : : .name = "itertools.count",
4165 : : .basicsize = sizeof(countobject),
4166 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
4167 : : Py_TPFLAGS_IMMUTABLETYPE),
4168 : : .slots = count_slots,
4169 : : };
4170 : :
4171 : :
4172 : : /* repeat object ************************************************************/
4173 : :
4174 : : typedef struct {
4175 : : PyObject_HEAD
4176 : : PyObject *element;
4177 : : Py_ssize_t cnt;
4178 : : } repeatobject;
4179 : :
4180 : : static PyObject *
4181 : 0 : repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4182 : : {
4183 : : repeatobject *ro;
4184 : : PyObject *element;
4185 : 0 : Py_ssize_t cnt = -1, n_args;
4186 : : static char *kwargs[] = {"object", "times", NULL};
4187 : :
4188 : 0 : n_args = PyTuple_GET_SIZE(args);
4189 [ # # ]: 0 : if (kwds != NULL)
4190 : 0 : n_args += PyDict_GET_SIZE(kwds);
4191 [ # # ]: 0 : if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
4192 : : &element, &cnt))
4193 : 0 : return NULL;
4194 : : /* Does user supply times argument? */
4195 [ # # # # ]: 0 : if (n_args == 2 && cnt < 0)
4196 : 0 : cnt = 0;
4197 : :
4198 : 0 : ro = (repeatobject *)type->tp_alloc(type, 0);
4199 [ # # ]: 0 : if (ro == NULL)
4200 : 0 : return NULL;
4201 : 0 : ro->element = Py_NewRef(element);
4202 : 0 : ro->cnt = cnt;
4203 : 0 : return (PyObject *)ro;
4204 : : }
4205 : :
4206 : : static void
4207 : 0 : repeat_dealloc(repeatobject *ro)
4208 : : {
4209 : 0 : PyTypeObject *tp = Py_TYPE(ro);
4210 : 0 : PyObject_GC_UnTrack(ro);
4211 : 0 : Py_XDECREF(ro->element);
4212 : 0 : tp->tp_free(ro);
4213 : 0 : Py_DECREF(tp);
4214 : 0 : }
4215 : :
4216 : : static int
4217 : 0 : repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
4218 : : {
4219 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(ro));
4220 [ # # # # ]: 0 : Py_VISIT(ro->element);
4221 : 0 : return 0;
4222 : : }
4223 : :
4224 : : static PyObject *
4225 : 0 : repeat_next(repeatobject *ro)
4226 : : {
4227 [ # # ]: 0 : if (ro->cnt == 0)
4228 : 0 : return NULL;
4229 [ # # ]: 0 : if (ro->cnt > 0)
4230 : 0 : ro->cnt--;
4231 : 0 : return Py_NewRef(ro->element);
4232 : : }
4233 : :
4234 : : static PyObject *
4235 : 0 : repeat_repr(repeatobject *ro)
4236 : : {
4237 [ # # ]: 0 : if (ro->cnt == -1)
4238 : 0 : return PyUnicode_FromFormat("%s(%R)",
4239 : : _PyType_Name(Py_TYPE(ro)), ro->element);
4240 : : else
4241 : 0 : return PyUnicode_FromFormat("%s(%R, %zd)",
4242 : : _PyType_Name(Py_TYPE(ro)), ro->element,
4243 : : ro->cnt);
4244 : : }
4245 : :
4246 : : static PyObject *
4247 : 0 : repeat_len(repeatobject *ro, PyObject *Py_UNUSED(ignored))
4248 : : {
4249 [ # # ]: 0 : if (ro->cnt == -1) {
4250 : 0 : PyErr_SetString(PyExc_TypeError, "len() of unsized object");
4251 : 0 : return NULL;
4252 : : }
4253 : 0 : return PyLong_FromSize_t(ro->cnt);
4254 : : }
4255 : :
4256 : : PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
4257 : :
4258 : : static PyObject *
4259 : 0 : repeat_reduce(repeatobject *ro, PyObject *Py_UNUSED(ignored))
4260 : : {
4261 : : /* unpickle this so that a new repeat iterator is constructed with an
4262 : : * object, then call __setstate__ on it to set cnt
4263 : : */
4264 [ # # ]: 0 : if (ro->cnt >= 0)
4265 : 0 : return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
4266 : : else
4267 : 0 : return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
4268 : : }
4269 : :
4270 : : static PyMethodDef repeat_methods[] = {
4271 : : {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
4272 : : {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
4273 : : {NULL, NULL} /* sentinel */
4274 : : };
4275 : :
4276 : : PyDoc_STRVAR(repeat_doc,
4277 : : "repeat(object [,times]) -> create an iterator which returns the object\n\
4278 : : for the specified number of times. If not specified, returns the object\n\
4279 : : endlessly.");
4280 : :
4281 : : static PyType_Slot repeat_slots[] = {
4282 : : {Py_tp_dealloc, repeat_dealloc},
4283 : : {Py_tp_repr, repeat_repr},
4284 : : {Py_tp_getattro, PyObject_GenericGetAttr},
4285 : : {Py_tp_doc, (void *)repeat_doc},
4286 : : {Py_tp_traverse, repeat_traverse},
4287 : : {Py_tp_iter, PyObject_SelfIter},
4288 : : {Py_tp_iternext, repeat_next},
4289 : : {Py_tp_methods, repeat_methods},
4290 : : {Py_tp_new, repeat_new},
4291 : : {Py_tp_free, PyObject_GC_Del},
4292 : : {0, NULL},
4293 : : };
4294 : :
4295 : : static PyType_Spec repeat_spec = {
4296 : : .name = "itertools.repeat",
4297 : : .basicsize = sizeof(repeatobject),
4298 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
4299 : : Py_TPFLAGS_IMMUTABLETYPE),
4300 : : .slots = repeat_slots,
4301 : : };
4302 : :
4303 : :
4304 : : /* ziplongest object *********************************************************/
4305 : :
4306 : : typedef struct {
4307 : : PyObject_HEAD
4308 : : Py_ssize_t tuplesize;
4309 : : Py_ssize_t numactive;
4310 : : PyObject *ittuple; /* tuple of iterators */
4311 : : PyObject *result;
4312 : : PyObject *fillvalue;
4313 : : } ziplongestobject;
4314 : :
4315 : : static PyObject *
4316 : 0 : zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
4317 : : {
4318 : : ziplongestobject *lz;
4319 : : Py_ssize_t i;
4320 : : PyObject *ittuple; /* tuple of iterators */
4321 : : PyObject *result;
4322 : 0 : PyObject *fillvalue = Py_None;
4323 : : Py_ssize_t tuplesize;
4324 : :
4325 [ # # # # : 0 : if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
# # ]
4326 : 0 : fillvalue = NULL;
4327 [ # # ]: 0 : if (PyDict_GET_SIZE(kwds) == 1) {
4328 : 0 : fillvalue = PyDict_GetItemWithError(kwds, &_Py_ID(fillvalue));
4329 : : }
4330 [ # # ]: 0 : if (fillvalue == NULL) {
4331 [ # # ]: 0 : if (!PyErr_Occurred()) {
4332 : 0 : PyErr_SetString(PyExc_TypeError,
4333 : : "zip_longest() got an unexpected keyword argument");
4334 : : }
4335 : 0 : return NULL;
4336 : : }
4337 : : }
4338 : :
4339 : : /* args must be a tuple */
4340 : : assert(PyTuple_Check(args));
4341 : 0 : tuplesize = PyTuple_GET_SIZE(args);
4342 : :
4343 : : /* obtain iterators */
4344 : 0 : ittuple = PyTuple_New(tuplesize);
4345 [ # # ]: 0 : if (ittuple == NULL)
4346 : 0 : return NULL;
4347 [ # # ]: 0 : for (i=0; i < tuplesize; i++) {
4348 : 0 : PyObject *item = PyTuple_GET_ITEM(args, i);
4349 : 0 : PyObject *it = PyObject_GetIter(item);
4350 [ # # ]: 0 : if (it == NULL) {
4351 : 0 : Py_DECREF(ittuple);
4352 : 0 : return NULL;
4353 : : }
4354 : 0 : PyTuple_SET_ITEM(ittuple, i, it);
4355 : : }
4356 : :
4357 : : /* create a result holder */
4358 : 0 : result = PyTuple_New(tuplesize);
4359 [ # # ]: 0 : if (result == NULL) {
4360 : 0 : Py_DECREF(ittuple);
4361 : 0 : return NULL;
4362 : : }
4363 [ # # ]: 0 : for (i=0 ; i < tuplesize ; i++) {
4364 : 0 : Py_INCREF(Py_None);
4365 : 0 : PyTuple_SET_ITEM(result, i, Py_None);
4366 : : }
4367 : :
4368 : : /* create ziplongestobject structure */
4369 : 0 : lz = (ziplongestobject *)type->tp_alloc(type, 0);
4370 [ # # ]: 0 : if (lz == NULL) {
4371 : 0 : Py_DECREF(ittuple);
4372 : 0 : Py_DECREF(result);
4373 : 0 : return NULL;
4374 : : }
4375 : 0 : lz->ittuple = ittuple;
4376 : 0 : lz->tuplesize = tuplesize;
4377 : 0 : lz->numactive = tuplesize;
4378 : 0 : lz->result = result;
4379 : 0 : lz->fillvalue = Py_NewRef(fillvalue);
4380 : 0 : return (PyObject *)lz;
4381 : : }
4382 : :
4383 : : static void
4384 : 0 : zip_longest_dealloc(ziplongestobject *lz)
4385 : : {
4386 : 0 : PyTypeObject *tp = Py_TYPE(lz);
4387 : 0 : PyObject_GC_UnTrack(lz);
4388 : 0 : Py_XDECREF(lz->ittuple);
4389 : 0 : Py_XDECREF(lz->result);
4390 : 0 : Py_XDECREF(lz->fillvalue);
4391 : 0 : tp->tp_free(lz);
4392 : 0 : Py_DECREF(tp);
4393 : 0 : }
4394 : :
4395 : : static int
4396 : 0 : zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
4397 : : {
4398 [ # # # # ]: 0 : Py_VISIT(Py_TYPE(lz));
4399 [ # # # # ]: 0 : Py_VISIT(lz->ittuple);
4400 [ # # # # ]: 0 : Py_VISIT(lz->result);
4401 [ # # # # ]: 0 : Py_VISIT(lz->fillvalue);
4402 : 0 : return 0;
4403 : : }
4404 : :
4405 : : static PyObject *
4406 : 0 : zip_longest_next(ziplongestobject *lz)
4407 : : {
4408 : : Py_ssize_t i;
4409 : 0 : Py_ssize_t tuplesize = lz->tuplesize;
4410 : 0 : PyObject *result = lz->result;
4411 : : PyObject *it;
4412 : : PyObject *item;
4413 : : PyObject *olditem;
4414 : :
4415 [ # # ]: 0 : if (tuplesize == 0)
4416 : 0 : return NULL;
4417 [ # # ]: 0 : if (lz->numactive == 0)
4418 : 0 : return NULL;
4419 [ # # ]: 0 : if (Py_REFCNT(result) == 1) {
4420 : 0 : Py_INCREF(result);
4421 [ # # ]: 0 : for (i=0 ; i < tuplesize ; i++) {
4422 : 0 : it = PyTuple_GET_ITEM(lz->ittuple, i);
4423 [ # # ]: 0 : if (it == NULL) {
4424 : 0 : item = Py_NewRef(lz->fillvalue);
4425 : : } else {
4426 : 0 : item = PyIter_Next(it);
4427 [ # # ]: 0 : if (item == NULL) {
4428 : 0 : lz->numactive -= 1;
4429 [ # # # # ]: 0 : if (lz->numactive == 0 || PyErr_Occurred()) {
4430 : 0 : lz->numactive = 0;
4431 : 0 : Py_DECREF(result);
4432 : 0 : return NULL;
4433 : : } else {
4434 : 0 : item = Py_NewRef(lz->fillvalue);
4435 : 0 : PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4436 : 0 : Py_DECREF(it);
4437 : : }
4438 : : }
4439 : : }
4440 : 0 : olditem = PyTuple_GET_ITEM(result, i);
4441 : 0 : PyTuple_SET_ITEM(result, i, item);
4442 : 0 : Py_DECREF(olditem);
4443 : : }
4444 : : // bpo-42536: The GC may have untracked this result tuple. Since we're
4445 : : // recycling it, make sure it's tracked again:
4446 [ # # ]: 0 : if (!_PyObject_GC_IS_TRACKED(result)) {
4447 : 0 : _PyObject_GC_TRACK(result);
4448 : : }
4449 : : } else {
4450 : 0 : result = PyTuple_New(tuplesize);
4451 [ # # ]: 0 : if (result == NULL)
4452 : 0 : return NULL;
4453 [ # # ]: 0 : for (i=0 ; i < tuplesize ; i++) {
4454 : 0 : it = PyTuple_GET_ITEM(lz->ittuple, i);
4455 [ # # ]: 0 : if (it == NULL) {
4456 : 0 : item = Py_NewRef(lz->fillvalue);
4457 : : } else {
4458 : 0 : item = PyIter_Next(it);
4459 [ # # ]: 0 : if (item == NULL) {
4460 : 0 : lz->numactive -= 1;
4461 [ # # # # ]: 0 : if (lz->numactive == 0 || PyErr_Occurred()) {
4462 : 0 : lz->numactive = 0;
4463 : 0 : Py_DECREF(result);
4464 : 0 : return NULL;
4465 : : } else {
4466 : 0 : item = Py_NewRef(lz->fillvalue);
4467 : 0 : PyTuple_SET_ITEM(lz->ittuple, i, NULL);
4468 : 0 : Py_DECREF(it);
4469 : : }
4470 : : }
4471 : : }
4472 : 0 : PyTuple_SET_ITEM(result, i, item);
4473 : : }
4474 : : }
4475 : 0 : return result;
4476 : : }
4477 : :
4478 : : static PyObject *
4479 : 0 : zip_longest_reduce(ziplongestobject *lz, PyObject *Py_UNUSED(ignored))
4480 : : {
4481 : :
4482 : : /* Create a new tuple with empty sequences where appropriate to pickle.
4483 : : * Then use setstate to set the fillvalue
4484 : : */
4485 : : int i;
4486 : 0 : PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
4487 : :
4488 [ # # ]: 0 : if (args == NULL)
4489 : 0 : return NULL;
4490 [ # # ]: 0 : for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
4491 : 0 : PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
4492 [ # # ]: 0 : if (elem == NULL) {
4493 : 0 : elem = PyTuple_New(0);
4494 [ # # ]: 0 : if (elem == NULL) {
4495 : 0 : Py_DECREF(args);
4496 : 0 : return NULL;
4497 : : }
4498 : : } else
4499 : 0 : Py_INCREF(elem);
4500 : 0 : PyTuple_SET_ITEM(args, i, elem);
4501 : : }
4502 : 0 : return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
4503 : : }
4504 : :
4505 : : static PyObject *
4506 : 0 : zip_longest_setstate(ziplongestobject *lz, PyObject *state)
4507 : : {
4508 : 0 : Py_INCREF(state);
4509 : 0 : Py_XSETREF(lz->fillvalue, state);
4510 : 0 : Py_RETURN_NONE;
4511 : : }
4512 : :
4513 : : static PyMethodDef zip_longest_methods[] = {
4514 : : {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
4515 : : reduce_doc},
4516 : : {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
4517 : : setstate_doc},
4518 : : {NULL, NULL} /* sentinel */
4519 : : };
4520 : :
4521 : : PyDoc_STRVAR(zip_longest_doc,
4522 : : "zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
4523 : : \n\
4524 : : Return a zip_longest object whose .__next__() method returns a tuple where\n\
4525 : : the i-th element comes from the i-th iterable argument. The .__next__()\n\
4526 : : method continues until the longest iterable in the argument sequence\n\
4527 : : is exhausted and then it raises StopIteration. When the shorter iterables\n\
4528 : : are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
4529 : : defaults to None or can be specified by a keyword argument.\n\
4530 : : ");
4531 : :
4532 : : static PyType_Slot ziplongest_slots[] = {
4533 : : {Py_tp_dealloc, zip_longest_dealloc},
4534 : : {Py_tp_getattro, PyObject_GenericGetAttr},
4535 : : {Py_tp_doc, (void *)zip_longest_doc},
4536 : : {Py_tp_traverse, zip_longest_traverse},
4537 : : {Py_tp_iter, PyObject_SelfIter},
4538 : : {Py_tp_iternext, zip_longest_next},
4539 : : {Py_tp_methods, zip_longest_methods},
4540 : : {Py_tp_new, zip_longest_new},
4541 : : {Py_tp_free, PyObject_GC_Del},
4542 : : {0, NULL},
4543 : : };
4544 : :
4545 : : static PyType_Spec ziplongest_spec = {
4546 : : .name = "itertools.zip_longest",
4547 : : .basicsize = sizeof(ziplongestobject),
4548 : : .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
4549 : : Py_TPFLAGS_IMMUTABLETYPE),
4550 : : .slots = ziplongest_slots,
4551 : : };
4552 : :
4553 : :
4554 : : /* module level code ********************************************************/
4555 : :
4556 : : PyDoc_STRVAR(module_doc,
4557 : : "Functional tools for creating and using iterators.\n\
4558 : : \n\
4559 : : Infinite iterators:\n\
4560 : : count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
4561 : : cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
4562 : : repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
4563 : : \n\
4564 : : Iterators terminating on the shortest input sequence:\n\
4565 : : accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
4566 : : batched(p, n) --> [p0, p1, ..., p_n-1], [p_n, p_n+1, ..., p_2n-1], ...\n\
4567 : : chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
4568 : : chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
4569 : : compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
4570 : : dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\
4571 : : groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
4572 : : filterfalse(pred, seq) --> elements of seq where pred(elem) is False\n\
4573 : : islice(seq, [start,] stop [, step]) --> elements from\n\
4574 : : seq[start:stop:step]\n\
4575 : : pairwise(s) --> (s[0],s[1]), (s[1],s[2]), (s[2], s[3]), ...\n\
4576 : : starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
4577 : : tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
4578 : : takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\
4579 : : zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
4580 : : \n\
4581 : : Combinatoric generators:\n\
4582 : : product(p, q, ... [repeat=1]) --> cartesian product\n\
4583 : : permutations(p[, r])\n\
4584 : : combinations(p, r)\n\
4585 : : combinations_with_replacement(p, r)\n\
4586 : : ");
4587 : :
4588 : : static int
4589 : 86 : itertoolsmodule_traverse(PyObject *mod, visitproc visit, void *arg)
4590 : : {
4591 : 86 : itertools_state *state = get_module_state(mod);
4592 [ + - - + ]: 86 : Py_VISIT(state->accumulate_type);
4593 [ + - - + ]: 86 : Py_VISIT(state->batched_type);
4594 [ + - - + ]: 86 : Py_VISIT(state->chain_type);
4595 [ + - - + ]: 86 : Py_VISIT(state->combinations_type);
4596 [ + - - + ]: 86 : Py_VISIT(state->compress_type);
4597 [ + - - + ]: 86 : Py_VISIT(state->count_type);
4598 [ + - - + ]: 86 : Py_VISIT(state->cwr_type);
4599 [ + - - + ]: 86 : Py_VISIT(state->cycle_type);
4600 [ + - - + ]: 86 : Py_VISIT(state->dropwhile_type);
4601 [ + - - + ]: 86 : Py_VISIT(state->filterfalse_type);
4602 [ + - - + ]: 86 : Py_VISIT(state->groupby_type);
4603 [ + - - + ]: 86 : Py_VISIT(state->_grouper_type);
4604 [ + - - + ]: 86 : Py_VISIT(state->islice_type);
4605 [ + - - + ]: 86 : Py_VISIT(state->pairwise_type);
4606 [ + - - + ]: 86 : Py_VISIT(state->permutations_type);
4607 [ + - - + ]: 86 : Py_VISIT(state->product_type);
4608 [ + - - + ]: 86 : Py_VISIT(state->repeat_type);
4609 [ + - - + ]: 86 : Py_VISIT(state->starmap_type);
4610 [ + - - + ]: 86 : Py_VISIT(state->takewhile_type);
4611 [ + - - + ]: 86 : Py_VISIT(state->tee_type);
4612 [ + - - + ]: 86 : Py_VISIT(state->teedataobject_type);
4613 [ + - - + ]: 86 : Py_VISIT(state->ziplongest_type);
4614 : 86 : return 0;
4615 : : }
4616 : :
4617 : : static int
4618 : 12 : itertoolsmodule_clear(PyObject *mod)
4619 : : {
4620 : 12 : itertools_state *state = get_module_state(mod);
4621 [ + + ]: 12 : Py_CLEAR(state->accumulate_type);
4622 [ + + ]: 12 : Py_CLEAR(state->batched_type);
4623 [ + + ]: 12 : Py_CLEAR(state->chain_type);
4624 [ + + ]: 12 : Py_CLEAR(state->combinations_type);
4625 [ + + ]: 12 : Py_CLEAR(state->compress_type);
4626 [ + + ]: 12 : Py_CLEAR(state->count_type);
4627 [ + + ]: 12 : Py_CLEAR(state->cwr_type);
4628 [ + + ]: 12 : Py_CLEAR(state->cycle_type);
4629 [ + + ]: 12 : Py_CLEAR(state->dropwhile_type);
4630 [ + + ]: 12 : Py_CLEAR(state->filterfalse_type);
4631 [ + + ]: 12 : Py_CLEAR(state->groupby_type);
4632 [ + + ]: 12 : Py_CLEAR(state->_grouper_type);
4633 [ + + ]: 12 : Py_CLEAR(state->islice_type);
4634 [ + + ]: 12 : Py_CLEAR(state->pairwise_type);
4635 [ + + ]: 12 : Py_CLEAR(state->permutations_type);
4636 [ + + ]: 12 : Py_CLEAR(state->product_type);
4637 [ + + ]: 12 : Py_CLEAR(state->repeat_type);
4638 [ + + ]: 12 : Py_CLEAR(state->starmap_type);
4639 [ + + ]: 12 : Py_CLEAR(state->takewhile_type);
4640 [ + + ]: 12 : Py_CLEAR(state->tee_type);
4641 [ + + ]: 12 : Py_CLEAR(state->teedataobject_type);
4642 [ + + ]: 12 : Py_CLEAR(state->ziplongest_type);
4643 : 12 : return 0;
4644 : : }
4645 : :
4646 : : static void
4647 : 6 : itertoolsmodule_free(void *mod)
4648 : : {
4649 : 6 : (void)itertoolsmodule_clear((PyObject *)mod);
4650 : 6 : }
4651 : :
4652 : : #define ADD_TYPE(module, type, spec) \
4653 : : do { \
4654 : : type = (PyTypeObject *)PyType_FromModuleAndSpec(module, spec, NULL); \
4655 : : if (type == NULL) { \
4656 : : return -1; \
4657 : : } \
4658 : : if (PyModule_AddType(module, type) < 0) { \
4659 : : return -1; \
4660 : : } \
4661 : : } while (0)
4662 : :
4663 : : static int
4664 : 6 : itertoolsmodule_exec(PyObject *mod)
4665 : : {
4666 : 6 : itertools_state *state = get_module_state(mod);
4667 [ - + - + ]: 6 : ADD_TYPE(mod, state->accumulate_type, &accumulate_spec);
4668 [ - + - + ]: 6 : ADD_TYPE(mod, state->batched_type, &batched_spec);
4669 [ - + - + ]: 6 : ADD_TYPE(mod, state->chain_type, &chain_spec);
4670 [ - + - + ]: 6 : ADD_TYPE(mod, state->combinations_type, &combinations_spec);
4671 [ - + - + ]: 6 : ADD_TYPE(mod, state->compress_type, &compress_spec);
4672 [ - + - + ]: 6 : ADD_TYPE(mod, state->count_type, &count_spec);
4673 [ - + - + ]: 6 : ADD_TYPE(mod, state->cwr_type, &cwr_spec);
4674 [ - + - + ]: 6 : ADD_TYPE(mod, state->cycle_type, &cycle_spec);
4675 [ - + - + ]: 6 : ADD_TYPE(mod, state->dropwhile_type, &dropwhile_spec);
4676 [ - + - + ]: 6 : ADD_TYPE(mod, state->filterfalse_type, &filterfalse_spec);
4677 [ - + - + ]: 6 : ADD_TYPE(mod, state->groupby_type, &groupby_spec);
4678 [ - + - + ]: 6 : ADD_TYPE(mod, state->_grouper_type, &_grouper_spec);
4679 [ - + - + ]: 6 : ADD_TYPE(mod, state->islice_type, &islice_spec);
4680 [ - + - + ]: 6 : ADD_TYPE(mod, state->pairwise_type, &pairwise_spec);
4681 [ - + - + ]: 6 : ADD_TYPE(mod, state->permutations_type, &permutations_spec);
4682 [ - + - + ]: 6 : ADD_TYPE(mod, state->product_type, &product_spec);
4683 [ - + - + ]: 6 : ADD_TYPE(mod, state->repeat_type, &repeat_spec);
4684 [ - + - + ]: 6 : ADD_TYPE(mod, state->starmap_type, &starmap_spec);
4685 [ - + - + ]: 6 : ADD_TYPE(mod, state->takewhile_type, &takewhile_spec);
4686 [ - + - + ]: 6 : ADD_TYPE(mod, state->tee_type, &tee_spec);
4687 [ - + - + ]: 6 : ADD_TYPE(mod, state->teedataobject_type, &teedataobject_spec);
4688 [ - + - + ]: 6 : ADD_TYPE(mod, state->ziplongest_type, &ziplongest_spec);
4689 : :
4690 : 6 : Py_SET_TYPE(state->teedataobject_type, &PyType_Type);
4691 : 6 : return 0;
4692 : : }
4693 : :
4694 : : static struct PyModuleDef_Slot itertoolsmodule_slots[] = {
4695 : : {Py_mod_exec, itertoolsmodule_exec},
4696 : : {0, NULL}
4697 : : };
4698 : :
4699 : : static PyMethodDef module_methods[] = {
4700 : : ITERTOOLS_TEE_METHODDEF
4701 : : {NULL, NULL} /* sentinel */
4702 : : };
4703 : :
4704 : :
4705 : : static struct PyModuleDef itertoolsmodule = {
4706 : : .m_base = PyModuleDef_HEAD_INIT,
4707 : : .m_name = "itertools",
4708 : : .m_doc = module_doc,
4709 : : .m_size = sizeof(itertools_state),
4710 : : .m_methods = module_methods,
4711 : : .m_slots = itertoolsmodule_slots,
4712 : : .m_traverse = itertoolsmodule_traverse,
4713 : : .m_clear = itertoolsmodule_clear,
4714 : : .m_free = itertoolsmodule_free,
4715 : : };
4716 : :
4717 : : PyMODINIT_FUNC
4718 : 6 : PyInit_itertools(void)
4719 : : {
4720 : 6 : return PyModuleDef_Init(&itertoolsmodule);
4721 : : }
|