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      350249 : PySeqIter_New(PyObject *seq)
      15             : {
      16             :     seqiterobject *it;
      17             : 
      18      350249 :     if (!PySequence_Check(seq)) {
      19           0 :         PyErr_BadInternalCall();
      20           0 :         return NULL;
      21             :     }
      22      350249 :     it = PyObject_GC_New(seqiterobject, &PySeqIter_Type);
      23      350249 :     if (it == NULL)
      24           0 :         return NULL;
      25      350249 :     it->it_index = 0;
      26      350249 :     Py_INCREF(seq);
      27      350249 :     it->it_seq = seq;
      28      350249 :     _PyObject_GC_TRACK(it);
      29      350249 :     return (PyObject *)it;
      30             : }
      31             : 
      32             : static void
      33      350249 : iter_dealloc(seqiterobject *it)
      34             : {
      35      350249 :     _PyObject_GC_UNTRACK(it);
      36      350249 :     Py_XDECREF(it->it_seq);
      37      350249 :     PyObject_GC_Del(it);
      38      350249 : }
      39             : 
      40             : static int
      41         385 : iter_traverse(seqiterobject *it, visitproc visit, void *arg)
      42             : {
      43         385 :     Py_VISIT(it->it_seq);
      44         385 :     return 0;
      45             : }
      46             : 
      47             : static PyObject *
      48     1123950 : iter_iternext(PyObject *iterator)
      49             : {
      50             :     seqiterobject *it;
      51             :     PyObject *seq;
      52             :     PyObject *result;
      53             : 
      54     1123950 :     assert(PySeqIter_Check(iterator));
      55     1123950 :     it = (seqiterobject *)iterator;
      56     1123950 :     seq = it->it_seq;
      57     1123950 :     if (seq == NULL)
      58          11 :         return NULL;
      59     1123940 :     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     1123930 :     result = PySequence_GetItem(seq, it->it_index);
      66     1123930 :     if (result != NULL) {
      67      774410 :         it->it_index++;
      68      774410 :         return result;
      69             :     }
      70      349542 :     if (PyErr_ExceptionMatches(PyExc_IndexError) ||
      71          19 :         PyErr_ExceptionMatches(PyExc_StopIteration))
      72             :     {
      73      349506 :         PyErr_Clear();
      74      349506 :         it->it_seq = NULL;
      75      349506 :         Py_DECREF(seq);
      76             :     }
      77      349523 :     return NULL;
      78             : }
      79             : 
      80             : static PyObject *
      81        9431 : iter_len(seqiterobject *it, PyObject *Py_UNUSED(ignored))
      82             : {
      83             :     Py_ssize_t seqsize, len;
      84             : 
      85        9431 :     if (it->it_seq) {
      86        9426 :         if (_PyObject_HasLen(it->it_seq)) {
      87        9379 :             seqsize = PySequence_Size(it->it_seq);
      88        9379 :             if (seqsize == -1)
      89           0 :                 return NULL;
      90             :         }
      91             :         else {
      92          47 :             Py_RETURN_NOTIMPLEMENTED;
      93             :         }
      94        9379 :         len = seqsize - it->it_index;
      95        9379 :         if (len >= 0)
      96        9379 :             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     3081640 : PyCallIter_New(PyObject *callable, PyObject *sentinel)
     181             : {
     182             :     calliterobject *it;
     183     3081640 :     it = PyObject_GC_New(calliterobject, &PyCallIter_Type);
     184     3081640 :     if (it == NULL)
     185           0 :         return NULL;
     186     3081640 :     Py_INCREF(callable);
     187     3081640 :     it->it_callable = callable;
     188     3081640 :     Py_INCREF(sentinel);
     189     3081640 :     it->it_sentinel = sentinel;
     190     3081640 :     _PyObject_GC_TRACK(it);
     191     3081640 :     return (PyObject *)it;
     192             : }
     193             : static void
     194     3081640 : calliter_dealloc(calliterobject *it)
     195             : {
     196     3081640 :     _PyObject_GC_UNTRACK(it);
     197     3081640 :     Py_XDECREF(it->it_callable);
     198     3081640 :     Py_XDECREF(it->it_sentinel);
     199     3081640 :     PyObject_GC_Del(it);
     200     3081640 : }
     201             : 
     202             : static int
     203        3458 : calliter_traverse(calliterobject *it, visitproc visit, void *arg)
     204             : {
     205        3458 :     Py_VISIT(it->it_callable);
     206        3458 :     Py_VISIT(it->it_sentinel);
     207        3458 :     return 0;
     208             : }
     209             : 
     210             : static PyObject *
     211     4097040 : calliter_iternext(calliterobject *it)
     212             : {
     213             :     PyObject *result;
     214             : 
     215     4097040 :     if (it->it_callable == NULL) {
     216          17 :         return NULL;
     217             :     }
     218             : 
     219     4097020 :     result = _PyObject_CallNoArgs(it->it_callable);
     220     4097020 :     if (result != NULL) {
     221             :         int ok;
     222             : 
     223     4094140 :         ok = PyObject_RichCompareBool(it->it_sentinel, result, Py_EQ);
     224     4094140 :         if (ok == 0) {
     225     1016060 :             return result; /* Common case, fast path */
     226             :         }
     227             : 
     228     3078080 :         Py_DECREF(result);
     229     3078080 :         if (ok > 0) {
     230     3078080 :             Py_CLEAR(it->it_callable);
     231     3078080 :             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     3080960 :     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          49 :     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           8 :         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             : }