1 | /* Iterator objects */ |
2 | |
3 | #include "Python.h" |
4 | #include "pycore_object.h" |
5 | |
6 | typedef struct { |
7 | PyObject_HEAD |
8 | Py_ssize_t it_index; |
9 | PyObject *it_seq; /* Set to NULL when iterator is exhausted */ |
10 | } seqiterobject; |
11 | |
12 | _Py_IDENTIFIER(iter); |
13 | |
14 | PyObject * |
15 | PySeqIter_New(PyObject *seq) |
16 | { |
17 | seqiterobject *it; |
18 | |
19 | if (!PySequence_Check(seq)) { |
20 | PyErr_BadInternalCall(); |
21 | return NULL; |
22 | } |
23 | it = PyObject_GC_New(seqiterobject, &PySeqIter_Type); |
24 | if (it == NULL) |
25 | return NULL; |
26 | it->it_index = 0; |
27 | Py_INCREF(seq); |
28 | it->it_seq = seq; |
29 | _PyObject_GC_TRACK(it); |
30 | return (PyObject *)it; |
31 | } |
32 | |
33 | static void |
34 | iter_dealloc(seqiterobject *it) |
35 | { |
36 | _PyObject_GC_UNTRACK(it); |
37 | Py_XDECREF(it->it_seq); |
38 | PyObject_GC_Del(it); |
39 | } |
40 | |
41 | static int |
42 | iter_traverse(seqiterobject *it, visitproc visit, void *arg) |
43 | { |
44 | Py_VISIT(it->it_seq); |
45 | return 0; |
46 | } |
47 | |
48 | static PyObject * |
49 | iter_iternext(PyObject *iterator) |
50 | { |
51 | seqiterobject *it; |
52 | PyObject *seq; |
53 | PyObject *result; |
54 | |
55 | assert(PySeqIter_Check(iterator)); |
56 | it = (seqiterobject *)iterator; |
57 | seq = it->it_seq; |
58 | if (seq == NULL) |
59 | return NULL; |
60 | if (it->it_index == PY_SSIZE_T_MAX) { |
61 | PyErr_SetString(PyExc_OverflowError, |
62 | "iter index too large" ); |
63 | return NULL; |
64 | } |
65 | |
66 | result = PySequence_GetItem(seq, it->it_index); |
67 | if (result != NULL) { |
68 | it->it_index++; |
69 | return result; |
70 | } |
71 | if (PyErr_ExceptionMatches(PyExc_IndexError) || |
72 | PyErr_ExceptionMatches(PyExc_StopIteration)) |
73 | { |
74 | PyErr_Clear(); |
75 | it->it_seq = NULL; |
76 | Py_DECREF(seq); |
77 | } |
78 | return NULL; |
79 | } |
80 | |
81 | static PyObject * |
82 | iter_len(seqiterobject *it, PyObject *Py_UNUSED(ignored)) |
83 | { |
84 | Py_ssize_t seqsize, len; |
85 | |
86 | if (it->it_seq) { |
87 | if (_PyObject_HasLen(it->it_seq)) { |
88 | seqsize = PySequence_Size(it->it_seq); |
89 | if (seqsize == -1) |
90 | return NULL; |
91 | } |
92 | else { |
93 | Py_RETURN_NOTIMPLEMENTED; |
94 | } |
95 | len = seqsize - it->it_index; |
96 | if (len >= 0) |
97 | return PyLong_FromSsize_t(len); |
98 | } |
99 | return PyLong_FromLong(0); |
100 | } |
101 | |
102 | PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it))." ); |
103 | |
104 | static PyObject * |
105 | iter_reduce(seqiterobject *it, PyObject *Py_UNUSED(ignored)) |
106 | { |
107 | if (it->it_seq != NULL) |
108 | return Py_BuildValue("N(O)n" , _PyEval_GetBuiltinId(&PyId_iter), |
109 | it->it_seq, it->it_index); |
110 | else |
111 | return Py_BuildValue("N(())" , _PyEval_GetBuiltinId(&PyId_iter)); |
112 | } |
113 | |
114 | PyDoc_STRVAR(reduce_doc, "Return state information for pickling." ); |
115 | |
116 | static PyObject * |
117 | iter_setstate(seqiterobject *it, PyObject *state) |
118 | { |
119 | Py_ssize_t index = PyLong_AsSsize_t(state); |
120 | if (index == -1 && PyErr_Occurred()) |
121 | return NULL; |
122 | if (it->it_seq != NULL) { |
123 | if (index < 0) |
124 | index = 0; |
125 | it->it_index = index; |
126 | } |
127 | Py_RETURN_NONE; |
128 | } |
129 | |
130 | PyDoc_STRVAR(setstate_doc, "Set state information for unpickling." ); |
131 | |
132 | static PyMethodDef seqiter_methods[] = { |
133 | {"__length_hint__" , (PyCFunction)iter_len, METH_NOARGS, length_hint_doc}, |
134 | {"__reduce__" , (PyCFunction)iter_reduce, METH_NOARGS, reduce_doc}, |
135 | {"__setstate__" , (PyCFunction)iter_setstate, METH_O, setstate_doc}, |
136 | {NULL, NULL} /* sentinel */ |
137 | }; |
138 | |
139 | PyTypeObject PySeqIter_Type = { |
140 | PyVarObject_HEAD_INIT(&PyType_Type, 0) |
141 | "iterator" , /* tp_name */ |
142 | sizeof(seqiterobject), /* tp_basicsize */ |
143 | 0, /* tp_itemsize */ |
144 | /* methods */ |
145 | (destructor)iter_dealloc, /* tp_dealloc */ |
146 | 0, /* tp_vectorcall_offset */ |
147 | 0, /* tp_getattr */ |
148 | 0, /* tp_setattr */ |
149 | 0, /* tp_as_async */ |
150 | 0, /* tp_repr */ |
151 | 0, /* tp_as_number */ |
152 | 0, /* tp_as_sequence */ |
153 | 0, /* tp_as_mapping */ |
154 | 0, /* tp_hash */ |
155 | 0, /* tp_call */ |
156 | 0, /* tp_str */ |
157 | PyObject_GenericGetAttr, /* tp_getattro */ |
158 | 0, /* tp_setattro */ |
159 | 0, /* tp_as_buffer */ |
160 | Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */ |
161 | 0, /* tp_doc */ |
162 | (traverseproc)iter_traverse, /* tp_traverse */ |
163 | 0, /* tp_clear */ |
164 | 0, /* tp_richcompare */ |
165 | 0, /* tp_weaklistoffset */ |
166 | PyObject_SelfIter, /* tp_iter */ |
167 | iter_iternext, /* tp_iternext */ |
168 | seqiter_methods, /* tp_methods */ |
169 | 0, /* tp_members */ |
170 | }; |
171 | |
172 | /* -------------------------------------- */ |
173 | |
174 | typedef struct { |
175 | PyObject_HEAD |
176 | PyObject *it_callable; /* Set to NULL when iterator is exhausted */ |
177 | PyObject *it_sentinel; /* Set to NULL when iterator is exhausted */ |
178 | } calliterobject; |
179 | |
180 | PyObject * |
181 | PyCallIter_New(PyObject *callable, PyObject *sentinel) |
182 | { |
183 | calliterobject *it; |
184 | it = PyObject_GC_New(calliterobject, &PyCallIter_Type); |
185 | if (it == NULL) |
186 | return NULL; |
187 | Py_INCREF(callable); |
188 | it->it_callable = callable; |
189 | Py_INCREF(sentinel); |
190 | it->it_sentinel = sentinel; |
191 | _PyObject_GC_TRACK(it); |
192 | return (PyObject *)it; |
193 | } |
194 | static void |
195 | calliter_dealloc(calliterobject *it) |
196 | { |
197 | _PyObject_GC_UNTRACK(it); |
198 | Py_XDECREF(it->it_callable); |
199 | Py_XDECREF(it->it_sentinel); |
200 | PyObject_GC_Del(it); |
201 | } |
202 | |
203 | static int |
204 | calliter_traverse(calliterobject *it, visitproc visit, void *arg) |
205 | { |
206 | Py_VISIT(it->it_callable); |
207 | Py_VISIT(it->it_sentinel); |
208 | return 0; |
209 | } |
210 | |
211 | static PyObject * |
212 | calliter_iternext(calliterobject *it) |
213 | { |
214 | PyObject *result; |
215 | |
216 | if (it->it_callable == NULL) { |
217 | return NULL; |
218 | } |
219 | |
220 | result = _PyObject_CallNoArg(it->it_callable); |
221 | if (result != NULL) { |
222 | int ok; |
223 | |
224 | ok = PyObject_RichCompareBool(it->it_sentinel, result, Py_EQ); |
225 | if (ok == 0) { |
226 | return result; /* Common case, fast path */ |
227 | } |
228 | |
229 | Py_DECREF(result); |
230 | if (ok > 0) { |
231 | Py_CLEAR(it->it_callable); |
232 | Py_CLEAR(it->it_sentinel); |
233 | } |
234 | } |
235 | else if (PyErr_ExceptionMatches(PyExc_StopIteration)) { |
236 | PyErr_Clear(); |
237 | Py_CLEAR(it->it_callable); |
238 | Py_CLEAR(it->it_sentinel); |
239 | } |
240 | return NULL; |
241 | } |
242 | |
243 | static PyObject * |
244 | calliter_reduce(calliterobject *it, PyObject *Py_UNUSED(ignored)) |
245 | { |
246 | if (it->it_callable != NULL && it->it_sentinel != NULL) |
247 | return Py_BuildValue("N(OO)" , _PyEval_GetBuiltinId(&PyId_iter), |
248 | it->it_callable, it->it_sentinel); |
249 | else |
250 | return Py_BuildValue("N(())" , _PyEval_GetBuiltinId(&PyId_iter)); |
251 | } |
252 | |
253 | static PyMethodDef calliter_methods[] = { |
254 | {"__reduce__" , (PyCFunction)calliter_reduce, METH_NOARGS, reduce_doc}, |
255 | {NULL, NULL} /* sentinel */ |
256 | }; |
257 | |
258 | PyTypeObject PyCallIter_Type = { |
259 | PyVarObject_HEAD_INIT(&PyType_Type, 0) |
260 | "callable_iterator" , /* tp_name */ |
261 | sizeof(calliterobject), /* tp_basicsize */ |
262 | 0, /* tp_itemsize */ |
263 | /* methods */ |
264 | (destructor)calliter_dealloc, /* tp_dealloc */ |
265 | 0, /* tp_vectorcall_offset */ |
266 | 0, /* tp_getattr */ |
267 | 0, /* tp_setattr */ |
268 | 0, /* tp_as_async */ |
269 | 0, /* tp_repr */ |
270 | 0, /* tp_as_number */ |
271 | 0, /* tp_as_sequence */ |
272 | 0, /* tp_as_mapping */ |
273 | 0, /* tp_hash */ |
274 | 0, /* tp_call */ |
275 | 0, /* tp_str */ |
276 | PyObject_GenericGetAttr, /* tp_getattro */ |
277 | 0, /* tp_setattro */ |
278 | 0, /* tp_as_buffer */ |
279 | Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */ |
280 | 0, /* tp_doc */ |
281 | (traverseproc)calliter_traverse, /* tp_traverse */ |
282 | 0, /* tp_clear */ |
283 | 0, /* tp_richcompare */ |
284 | 0, /* tp_weaklistoffset */ |
285 | PyObject_SelfIter, /* tp_iter */ |
286 | (iternextfunc)calliter_iternext, /* tp_iternext */ |
287 | calliter_methods, /* tp_methods */ |
288 | }; |
289 | |
290 | |
291 | /* -------------------------------------- */ |
292 | |
293 | typedef struct { |
294 | PyObject_HEAD |
295 | PyObject *wrapped; |
296 | PyObject *default_value; |
297 | } anextawaitableobject; |
298 | |
299 | static void |
300 | anextawaitable_dealloc(anextawaitableobject *obj) |
301 | { |
302 | _PyObject_GC_UNTRACK(obj); |
303 | Py_XDECREF(obj->wrapped); |
304 | Py_XDECREF(obj->default_value); |
305 | PyObject_GC_Del(obj); |
306 | } |
307 | |
308 | static int |
309 | anextawaitable_traverse(anextawaitableobject *obj, visitproc visit, void *arg) |
310 | { |
311 | Py_VISIT(obj->wrapped); |
312 | Py_VISIT(obj->default_value); |
313 | return 0; |
314 | } |
315 | |
316 | static PyObject * |
317 | anextawaitable_getiter(anextawaitableobject *obj) |
318 | { |
319 | assert(obj->wrapped != NULL); |
320 | PyObject *awaitable = _PyCoro_GetAwaitableIter(obj->wrapped); |
321 | if (awaitable == NULL) { |
322 | return NULL; |
323 | } |
324 | if (Py_TYPE(awaitable)->tp_iternext == NULL) { |
325 | /* _PyCoro_GetAwaitableIter returns a Coroutine, a Generator, |
326 | * or an iterator. Of these, only coroutines lack tp_iternext. |
327 | */ |
328 | assert(PyCoro_CheckExact(awaitable)); |
329 | unaryfunc getter = Py_TYPE(awaitable)->tp_as_async->am_await; |
330 | PyObject *new_awaitable = getter(awaitable); |
331 | if (new_awaitable == NULL) { |
332 | Py_DECREF(awaitable); |
333 | return NULL; |
334 | } |
335 | Py_SETREF(awaitable, new_awaitable); |
336 | if (!PyIter_Check(awaitable)) { |
337 | PyErr_SetString(PyExc_TypeError, |
338 | "__await__ returned a non-iterable" ); |
339 | Py_DECREF(awaitable); |
340 | return NULL; |
341 | } |
342 | } |
343 | return awaitable; |
344 | } |
345 | |
346 | static PyObject * |
347 | anextawaitable_iternext(anextawaitableobject *obj) |
348 | { |
349 | /* Consider the following class: |
350 | * |
351 | * class A: |
352 | * async def __anext__(self): |
353 | * ... |
354 | * a = A() |
355 | * |
356 | * Then `await anext(a)` should call |
357 | * a.__anext__().__await__().__next__() |
358 | * |
359 | * On the other hand, given |
360 | * |
361 | * async def agen(): |
362 | * yield 1 |
363 | * yield 2 |
364 | * gen = agen() |
365 | * |
366 | * Then `await anext(gen)` can just call |
367 | * gen.__anext__().__next__() |
368 | */ |
369 | PyObject *awaitable = anextawaitable_getiter(obj); |
370 | if (awaitable == NULL) { |
371 | return NULL; |
372 | } |
373 | PyObject *result = (*Py_TYPE(awaitable)->tp_iternext)(awaitable); |
374 | Py_DECREF(awaitable); |
375 | if (result != NULL) { |
376 | return result; |
377 | } |
378 | if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration)) { |
379 | _PyGen_SetStopIterationValue(obj->default_value); |
380 | } |
381 | return NULL; |
382 | } |
383 | |
384 | |
385 | static PyObject * |
386 | anextawaitable_proxy(anextawaitableobject *obj, char *meth, PyObject *arg) { |
387 | PyObject *awaitable = anextawaitable_getiter(obj); |
388 | if (awaitable == NULL) { |
389 | return NULL; |
390 | } |
391 | PyObject *ret = PyObject_CallMethod(awaitable, meth, "O" , arg); |
392 | Py_DECREF(awaitable); |
393 | if (ret != NULL) { |
394 | return ret; |
395 | } |
396 | if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration)) { |
397 | /* `anextawaitableobject` is only used by `anext()` when |
398 | * a default value is provided. So when we have a StopAsyncIteration |
399 | * exception we replace it with a `StopIteration(default)`, as if |
400 | * it was the return value of `__anext__()` coroutine. |
401 | */ |
402 | _PyGen_SetStopIterationValue(obj->default_value); |
403 | } |
404 | return NULL; |
405 | } |
406 | |
407 | |
408 | static PyObject * |
409 | anextawaitable_send(anextawaitableobject *obj, PyObject *arg) { |
410 | return anextawaitable_proxy(obj, "send" , arg); |
411 | } |
412 | |
413 | |
414 | static PyObject * |
415 | anextawaitable_throw(anextawaitableobject *obj, PyObject *arg) { |
416 | return anextawaitable_proxy(obj, "throw" , arg); |
417 | } |
418 | |
419 | |
420 | static PyObject * |
421 | anextawaitable_close(anextawaitableobject *obj, PyObject *arg) { |
422 | return anextawaitable_proxy(obj, "close" , arg); |
423 | } |
424 | |
425 | |
426 | PyDoc_STRVAR(send_doc, |
427 | "send(arg) -> send 'arg' into the wrapped iterator,\n\ |
428 | return next yielded value or raise StopIteration." ); |
429 | |
430 | |
431 | PyDoc_STRVAR(throw_doc, |
432 | "throw(typ[,val[,tb]]) -> raise exception in the wrapped iterator,\n\ |
433 | return next yielded value or raise StopIteration." ); |
434 | |
435 | |
436 | PyDoc_STRVAR(close_doc, |
437 | "close() -> raise GeneratorExit inside generator." ); |
438 | |
439 | |
440 | static PyMethodDef anextawaitable_methods[] = { |
441 | {"send" ,(PyCFunction)anextawaitable_send, METH_O, send_doc}, |
442 | {"throw" ,(PyCFunction)anextawaitable_throw, METH_VARARGS, throw_doc}, |
443 | {"close" ,(PyCFunction)anextawaitable_close, METH_VARARGS, close_doc}, |
444 | {NULL, NULL} /* Sentinel */ |
445 | }; |
446 | |
447 | |
448 | static PyAsyncMethods anextawaitable_as_async = { |
449 | PyObject_SelfIter, /* am_await */ |
450 | 0, /* am_aiter */ |
451 | 0, /* am_anext */ |
452 | 0, /* am_send */ |
453 | }; |
454 | |
455 | PyTypeObject _PyAnextAwaitable_Type = { |
456 | PyVarObject_HEAD_INIT(&PyType_Type, 0) |
457 | "anext_awaitable" , /* tp_name */ |
458 | sizeof(anextawaitableobject), /* tp_basicsize */ |
459 | 0, /* tp_itemsize */ |
460 | /* methods */ |
461 | (destructor)anextawaitable_dealloc, /* tp_dealloc */ |
462 | 0, /* tp_vectorcall_offset */ |
463 | 0, /* tp_getattr */ |
464 | 0, /* tp_setattr */ |
465 | &anextawaitable_as_async, /* tp_as_async */ |
466 | 0, /* tp_repr */ |
467 | 0, /* tp_as_number */ |
468 | 0, /* tp_as_sequence */ |
469 | 0, /* tp_as_mapping */ |
470 | 0, /* tp_hash */ |
471 | 0, /* tp_call */ |
472 | 0, /* tp_str */ |
473 | PyObject_GenericGetAttr, /* tp_getattro */ |
474 | 0, /* tp_setattro */ |
475 | 0, /* tp_as_buffer */ |
476 | Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */ |
477 | 0, /* tp_doc */ |
478 | (traverseproc)anextawaitable_traverse, /* tp_traverse */ |
479 | 0, /* tp_clear */ |
480 | 0, /* tp_richcompare */ |
481 | 0, /* tp_weaklistoffset */ |
482 | PyObject_SelfIter, /* tp_iter */ |
483 | (unaryfunc)anextawaitable_iternext, /* tp_iternext */ |
484 | anextawaitable_methods, /* tp_methods */ |
485 | }; |
486 | |
487 | PyObject * |
488 | PyAnextAwaitable_New(PyObject *awaitable, PyObject *default_value) |
489 | { |
490 | anextawaitableobject *anext = PyObject_GC_New( |
491 | anextawaitableobject, &_PyAnextAwaitable_Type); |
492 | if (anext == NULL) { |
493 | return NULL; |
494 | } |
495 | Py_INCREF(awaitable); |
496 | anext->wrapped = awaitable; |
497 | Py_INCREF(default_value); |
498 | anext->default_value = default_value; |
499 | _PyObject_GC_TRACK(anext); |
500 | return (PyObject *)anext; |
501 | } |
502 | |