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