Issue #28003: Implement PEP 525 -- Asynchronous Generators.

6 files changed, 220 insertions, 36 deletions

diff --git a/Python/ceval.c b/Python/ceval.c
index a52ee8a..f737a2f 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -1204,7 +1204,7 @@ _PyEval_EvalFrameDefault(PyFrameObject *f, int throwflag)
     f->f_stacktop = NULL;       /* remains NULL unless yield suspends frame */
     f->f_executing = 1;
 
-    if (co->co_flags & (CO_GENERATOR | CO_COROUTINE)) {
+    if (co->co_flags & (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR)) {
         if (!throwflag && f->f_exc_type != NULL && f->f_exc_type != Py_None) {
             /* We were in an except handler when we left,
                restore the exception state which was put aside
@@ -2083,36 +2083,45 @@ _PyEval_EvalFrameDefault(PyFrameObject *f, int throwflag)
             PyObject *aiter = TOP();
             PyTypeObject *type = Py_TYPE(aiter);
 
-            if (type->tp_as_async != NULL)
-                getter = type->tp_as_async->am_anext;
+            if (PyAsyncGen_CheckExact(aiter)) {
+                awaitable = type->tp_as_async->am_anext(aiter);
+                if (awaitable == NULL) {
+                    goto error;
+                }
+            } else {
+                if (type->tp_as_async != NULL){
+                    getter = type->tp_as_async->am_anext;
+                }
 
-            if (getter != NULL) {
-                next_iter = (*getter)(aiter);
-                if (next_iter == NULL) {
+                if (getter != NULL) {
+                    next_iter = (*getter)(aiter);
+                    if (next_iter == NULL) {
+                        goto error;
+                    }
+                }
+                else {
+                    PyErr_Format(
+                        PyExc_TypeError,
+                        "'async for' requires an iterator with "
+                        "__anext__ method, got %.100s",
+                        type->tp_name);
                     goto error;
                 }
-            }
-            else {
-                PyErr_Format(
-                    PyExc_TypeError,
-                    "'async for' requires an iterator with "
-                    "__anext__ method, got %.100s",
-                    type->tp_name);
-                goto error;
-            }
 
-            awaitable = _PyCoro_GetAwaitableIter(next_iter);
-            if (awaitable == NULL) {
-                PyErr_Format(
-                    PyExc_TypeError,
-                    "'async for' received an invalid object "
-                    "from __anext__: %.100s",
-                    Py_TYPE(next_iter)->tp_name);
+                awaitable = _PyCoro_GetAwaitableIter(next_iter);
+                if (awaitable == NULL) {
+                    PyErr_Format(
+                        PyExc_TypeError,
+                        "'async for' received an invalid object "
+                        "from __anext__: %.100s",
+                        Py_TYPE(next_iter)->tp_name);
 
-                Py_DECREF(next_iter);
-                goto error;
-            } else
-                Py_DECREF(next_iter);
+                    Py_DECREF(next_iter);
+                    goto error;
+                } else {
+                    Py_DECREF(next_iter);
+                }
+            }
 
             PUSH(awaitable);
             PREDICT(LOAD_CONST);
@@ -2187,6 +2196,17 @@ _PyEval_EvalFrameDefault(PyFrameObject *f, int throwflag)
 
         TARGET(YIELD_VALUE) {
             retval = POP();
+
+            if (co->co_flags & CO_ASYNC_GENERATOR) {
+                PyObject *w = _PyAsyncGenValueWrapperNew(retval);
+                Py_DECREF(retval);
+                if (w == NULL) {
+                    retval = NULL;
+                    goto error;
+                }
+                retval = w;
+            }
+
             f->f_stacktop = stack_pointer;
             why = WHY_YIELD;
             goto fast_yield;
@@ -3712,7 +3732,7 @@ fast_block_end:
     assert((retval != NULL) ^ (PyErr_Occurred() != NULL));
 
 fast_yield:
-    if (co->co_flags & (CO_GENERATOR | CO_COROUTINE)) {
+    if (co->co_flags & (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR)) {
 
         /* The purpose of this block is to put aside the generator's exception
            state and restore that of the calling frame. If the current
@@ -4156,8 +4176,8 @@ _PyEval_EvalCodeWithName(PyObject *_co, PyObject *globals, PyObject *locals,
         freevars[PyTuple_GET_SIZE(co->co_cellvars) + i] = o;
     }
 
-    /* Handle generator/coroutine */
-    if (co->co_flags & (CO_GENERATOR | CO_COROUTINE)) {
+    /* Handle generator/coroutine/asynchronous generator */
+    if (co->co_flags & (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR)) {
         PyObject *gen;
         PyObject *coro_wrapper = tstate->coroutine_wrapper;
         int is_coro = co->co_flags & CO_COROUTINE;
@@ -4182,6 +4202,8 @@ _PyEval_EvalCodeWithName(PyObject *_co, PyObject *globals, PyObject *locals,
          * and return that as the value. */
         if (is_coro) {
             gen = PyCoro_New(f, name, qualname);
+        } else if (co->co_flags & CO_ASYNC_GENERATOR) {
+            gen = PyAsyncGen_New(f, name, qualname);
         } else {
             gen = PyGen_NewWithQualName(f, name, qualname);
         }
@@ -4660,6 +4682,38 @@ _PyEval_GetCoroutineWrapper(void)
     return tstate->coroutine_wrapper;
 }
 
+void
+_PyEval_SetAsyncGenFirstiter(PyObject *firstiter)
+{
+    PyThreadState *tstate = PyThreadState_GET();
+
+    Py_XINCREF(firstiter);
+    Py_XSETREF(tstate->async_gen_firstiter, firstiter);
+}
+
+PyObject *
+_PyEval_GetAsyncGenFirstiter(void)
+{
+    PyThreadState *tstate = PyThreadState_GET();
+    return tstate->async_gen_firstiter;
+}
+
+void
+_PyEval_SetAsyncGenFinalizer(PyObject *finalizer)
+{
+    PyThreadState *tstate = PyThreadState_GET();
+
+    Py_XINCREF(finalizer);
+    Py_XSETREF(tstate->async_gen_finalizer, finalizer);
+}
+
+PyObject *
+_PyEval_GetAsyncGenFinalizer(void)
+{
+    PyThreadState *tstate = PyThreadState_GET();
+    return tstate->async_gen_finalizer;
+}
+
 PyObject *
 PyEval_GetBuiltins(void)
 {
diff --git a/Python/compile.c b/Python/compile.c
index b46edd4..faae4f5 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -1886,8 +1886,6 @@ compiler_function(struct compiler *c, stmt_ty s, int is_async)
         return 0;
     }
 
-    if (is_async)
-        co->co_flags |= CO_COROUTINE;
     compiler_make_closure(c, co, funcflags, qualname);
     Py_DECREF(qualname);
     Py_DECREF(co);
@@ -2801,6 +2799,9 @@ compiler_visit_stmt(struct compiler *c, stmt_ty s)
         if (c->u->u_ste->ste_type != FunctionBlock)
             return compiler_error(c, "'return' outside function");
         if (s->v.Return.value) {
+            if (c->u->u_ste->ste_coroutine && c->u->u_ste->ste_generator)
+                return compiler_error(
+                    c, "'return' with value in async generator");
             VISIT(c, expr, s->v.Return.value);
         }
         else
@@ -4115,8 +4116,6 @@ compiler_visit_expr(struct compiler *c, expr_ty e)
     case Yield_kind:
         if (c->u->u_ste->ste_type != FunctionBlock)
             return compiler_error(c, "'yield' outside function");
-        if (c->u->u_scope_type == COMPILER_SCOPE_ASYNC_FUNCTION)
-            return compiler_error(c, "'yield' inside async function");
         if (e->v.Yield.value) {
             VISIT(c, expr, e->v.Yield.value);
         }
@@ -4992,8 +4991,12 @@ compute_code_flags(struct compiler *c)
         flags |= CO_NEWLOCALS | CO_OPTIMIZED;
         if (ste->ste_nested)
             flags |= CO_NESTED;
-        if (ste->ste_generator)
+        if (ste->ste_generator && !ste->ste_coroutine)
             flags |= CO_GENERATOR;
+        if (!ste->ste_generator && ste->ste_coroutine)
+            flags |= CO_COROUTINE;
+        if (ste->ste_generator && ste->ste_coroutine)
+            flags |= CO_ASYNC_GENERATOR;
         if (ste->ste_varargs)
             flags |= CO_VARARGS;
         if (ste->ste_varkeywords)
diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c
index a2399ed..f93afd2 100644
--- a/Python/pylifecycle.c
+++ b/Python/pylifecycle.c
@@ -694,6 +694,7 @@ Py_FinalizeEx(void)
     _PyGC_Fini();
     _PyRandom_Fini();
     _PyArg_Fini();
+    PyAsyncGen_Fini();
 
     /* Cleanup Unicode implementation */
     _PyUnicode_Fini();
diff --git a/Python/pystate.c b/Python/pystate.c
index 959354d..a0a8c97 100644
--- a/Python/pystate.c
+++ b/Python/pystate.c
@@ -229,6 +229,9 @@ new_threadstate(PyInterpreterState *interp, int init)
         tstate->in_coroutine_wrapper = 0;
         tstate->co_extra_user_count = 0;
 
+        tstate->async_gen_firstiter = NULL;
+        tstate->async_gen_finalizer = NULL;
+
         if (init)
             _PyThreadState_Init(tstate);
 
@@ -408,6 +411,8 @@ PyThreadState_Clear(PyThreadState *tstate)
     Py_CLEAR(tstate->c_traceobj);
 
     Py_CLEAR(tstate->coroutine_wrapper);
+    Py_CLEAR(tstate->async_gen_firstiter);
+    Py_CLEAR(tstate->async_gen_finalizer);
 }
 
 
diff --git a/Python/symtable.c b/Python/symtable.c
index b8d9398..e55813f 100644
--- a/Python/symtable.c
+++ b/Python/symtable.c
@@ -63,6 +63,7 @@ ste_new(struct symtable *st, identifier name, _Py_block_ty block,
         ste->ste_nested = 1;
     ste->ste_child_free = 0;
     ste->ste_generator = 0;
+    ste->ste_coroutine = 0;
     ste->ste_returns_value = 0;
     ste->ste_needs_class_closure = 0;
 
@@ -378,7 +379,7 @@ error_at_directive(PySTEntryObject *ste, PyObject *name)
                                        PyLong_AsLong(PyTuple_GET_ITEM(data, 2)));
 
             return 0;
-        }   
+        }
     }
     PyErr_SetString(PyExc_RuntimeError,
                     "BUG: internal directive bookkeeping broken");
@@ -1397,6 +1398,7 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
                                   FunctionBlock, (void *)s, s->lineno,
                                   s->col_offset))
             VISIT_QUIT(st, 0);
+        st->st_cur->ste_coroutine = 1;
         VISIT(st, arguments, s->v.AsyncFunctionDef.args);
         VISIT_SEQ(st, stmt, s->v.AsyncFunctionDef.body);
         if (!symtable_exit_block(st, s))
@@ -1492,7 +1494,7 @@ symtable_visit_expr(struct symtable *st, expr_ty e)
         break;
     case Await_kind:
         VISIT(st, expr, e->v.Await.value);
-        st->st_cur->ste_generator = 1;
+        st->st_cur->ste_coroutine = 1;
         break;
     case Compare_kind:
         VISIT(st, expr, e->v.Compare.left);
diff --git a/Python/sysmodule.c b/Python/sysmodule.c
index 0fe76b7..3a02ae9 100644
--- a/Python/sysmodule.c
+++ b/Python/sysmodule.c
@@ -717,6 +717,113 @@ Return the wrapper for coroutine objects set by sys.set_coroutine_wrapper."
 );
 
 
+static PyTypeObject AsyncGenHooksType;
+
+PyDoc_STRVAR(asyncgen_hooks_doc,
+"asyncgen_hooks\n\
+\n\
+A struct sequence providing information about asynhronous\n\
+generators hooks.  The attributes are read only.");
+
+static PyStructSequence_Field asyncgen_hooks_fields[] = {
+    {"firstiter", "Hook to intercept first iteration"},
+    {"finalizer", "Hook to intercept finalization"},
+    {0}
+};
+
+static PyStructSequence_Desc asyncgen_hooks_desc = {
+    "asyncgen_hooks",          /* name */
+    asyncgen_hooks_doc,        /* doc */
+    asyncgen_hooks_fields ,    /* fields */
+    2
+};
+
+
+static PyObject *
+sys_set_asyncgen_hooks(PyObject *self, PyObject *args, PyObject *kw)
+{
+    static char *keywords[] = {"firstiter", "finalizer", NULL};
+    PyObject *firstiter = NULL;
+    PyObject *finalizer = NULL;
+
+    if (!PyArg_ParseTupleAndKeywords(
+            args, kw, "|OO", keywords,
+            &firstiter, &finalizer)) {
+        return NULL;
+    }
+
+    if (finalizer && finalizer != Py_None) {
+        if (!PyCallable_Check(finalizer)) {
+            PyErr_Format(PyExc_TypeError,
+                         "callable finalizer expected, got %.50s",
+                         Py_TYPE(finalizer)->tp_name);
+            return NULL;
+        }
+        _PyEval_SetAsyncGenFinalizer(finalizer);
+    }
+    else if (finalizer == Py_None) {
+        _PyEval_SetAsyncGenFinalizer(NULL);
+    }
+
+    if (firstiter && firstiter != Py_None) {
+        if (!PyCallable_Check(firstiter)) {
+            PyErr_Format(PyExc_TypeError,
+                         "callable firstiter expected, got %.50s",
+                         Py_TYPE(firstiter)->tp_name);
+            return NULL;
+        }
+        _PyEval_SetAsyncGenFirstiter(firstiter);
+    }
+    else if (firstiter == Py_None) {
+        _PyEval_SetAsyncGenFirstiter(NULL);
+    }
+
+    Py_RETURN_NONE;
+}
+
+PyDoc_STRVAR(set_asyncgen_hooks_doc,
+"set_asyncgen_hooks(*, firstiter=None, finalizer=None)\n\
+\n\
+Set a finalizer for async generators objects."
+);
+
+static PyObject *
+sys_get_asyncgen_hooks(PyObject *self, PyObject *args)
+{
+    PyObject *res;
+    PyObject *firstiter = _PyEval_GetAsyncGenFirstiter();
+    PyObject *finalizer = _PyEval_GetAsyncGenFinalizer();
+
+    res = PyStructSequence_New(&AsyncGenHooksType);
+    if (res == NULL) {
+        return NULL;
+    }
+
+    if (firstiter == NULL) {
+        firstiter = Py_None;
+    }
+
+    if (finalizer == NULL) {
+        finalizer = Py_None;
+    }
+
+    Py_INCREF(firstiter);
+    PyStructSequence_SET_ITEM(res, 0, firstiter);
+
+    Py_INCREF(finalizer);
+    PyStructSequence_SET_ITEM(res, 1, finalizer);
+
+    return res;
+}
+
+PyDoc_STRVAR(get_asyncgen_hooks_doc,
+"get_asyncgen_hooks()\n\
+\n\
+Return a namedtuple of installed asynchronous generators hooks \
+(firstiter, finalizer)."
+);
+
+
 static PyTypeObject Hash_InfoType;
 
 PyDoc_STRVAR(hash_info_doc,
@@ -1315,6 +1422,10 @@ static PyMethodDef sys_methods[] = {
      set_coroutine_wrapper_doc},
     {"get_coroutine_wrapper", sys_get_coroutine_wrapper, METH_NOARGS,
      get_coroutine_wrapper_doc},
+    {"set_asyncgen_hooks", sys_set_asyncgen_hooks,
+     METH_VARARGS | METH_KEYWORDS, set_asyncgen_hooks_doc},
+    {"get_asyncgen_hooks", sys_get_asyncgen_hooks, METH_NOARGS,
+     get_asyncgen_hooks_doc},
     {NULL,              NULL}           /* sentinel */
 };
 
@@ -1950,6 +2061,14 @@ _PySys_Init(void)
     SET_SYS_FROM_STRING("thread_info", PyThread_GetInfo());
 #endif
 
+    /* initialize asyncgen_hooks */
+    if (AsyncGenHooksType.tp_name == NULL) {
+        if (PyStructSequence_InitType2(
+                &AsyncGenHooksType, &asyncgen_hooks_desc) < 0) {
+            return NULL;
+        }
+    }
+
 #undef SET_SYS_FROM_STRING
 #undef SET_SYS_FROM_STRING_BORROW
     if (PyErr_Occurred())