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