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