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