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