1 files changed, 223 insertions, 8 deletions

diff --git a/Python/ceval.c b/Python/ceval.c
index 5ee9077..8acd082 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -1926,11 +1926,133 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
             goto fast_block_end;
         }
 
+        TARGET(GET_AITER) {
+            getaiterfunc getter = NULL;
+            PyObject *iter = NULL;
+            PyObject *awaitable = NULL;
+            PyObject *obj = TOP();
+            PyTypeObject *type = Py_TYPE(obj);
+
+            if (type->tp_as_async != NULL)
+                getter = type->tp_as_async->am_aiter;
+
+            if (getter != NULL) {
+                iter = (*getter)(obj);
+                Py_DECREF(obj);
+                if (iter == NULL) {
+                    SET_TOP(NULL);
+                    goto error;
+                }
+            }
+            else {
+                SET_TOP(NULL);
+                PyErr_Format(
+                    PyExc_TypeError,
+                    "'async for' requires an object with "
+                    "__aiter__ method, got %.100s",
+                    type->tp_name);
+                Py_DECREF(obj);
+                goto error;
+            }
+
+            awaitable = _PyGen_GetAwaitableIter(iter);
+            if (awaitable == NULL) {
+                SET_TOP(NULL);
+                PyErr_Format(
+                    PyExc_TypeError,
+                    "'async for' received an invalid object "
+                    "from __aiter__: %.100s",
+                    Py_TYPE(iter)->tp_name);
+
+                Py_DECREF(iter);
+                goto error;
+            } else
+                Py_DECREF(iter);
+
+            SET_TOP(awaitable);
+            DISPATCH();
+        }
+
+        TARGET(GET_ANEXT) {
+            aiternextfunc getter = NULL;
+            PyObject *next_iter = NULL;
+            PyObject *awaitable = NULL;
+            PyObject *aiter = TOP();
+            PyTypeObject *type = Py_TYPE(aiter);
+
+            if (type->tp_as_async != NULL)
+                getter = type->tp_as_async->am_anext;
+
+            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;
+            }
+
+            awaitable = _PyGen_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);
+
+            PUSH(awaitable);
+            DISPATCH();
+        }
+
+        TARGET(GET_AWAITABLE) {
+            PyObject *iterable = TOP();
+            PyObject *iter = _PyGen_GetAwaitableIter(iterable);
+
+            Py_DECREF(iterable);
+
+            SET_TOP(iter); /* Even if it's NULL */
+
+            if (iter == NULL) {
+                goto error;
+            }
+
+            DISPATCH();
+        }
+
         TARGET(YIELD_FROM) {
             PyObject *v = POP();
             PyObject *reciever = TOP();
             int err;
             if (PyGen_CheckExact(reciever)) {
+                if (
+                    (((PyCodeObject*) \
+                        ((PyGenObject*)reciever)->gi_code)->co_flags &
+                                                        CO_COROUTINE)
+                    && !(co->co_flags & (CO_COROUTINE | CO_ITERABLE_COROUTINE)))
+                {
+                    /* If we're yielding-from a coroutine object from a regular
+                       generator object - raise an error. */
+
+                    Py_CLEAR(v);
+                    Py_CLEAR(reciever);
+                    SET_TOP(NULL);
+
+                    PyErr_SetString(PyExc_TypeError,
+                                    "cannot 'yield from' a coroutine object "
+                                    "from a generator");
+                    goto error;
+                }
                 retval = _PyGen_Send((PyGenObject *)reciever, v);
             } else {
                 _Py_IDENTIFIER(send);
@@ -2822,11 +2944,26 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
         TARGET(GET_ITER) {
             /* before: [obj]; after [getiter(obj)] */
             PyObject *iterable = TOP();
-            PyObject *iter = PyObject_GetIter(iterable);
-            Py_DECREF(iterable);
-            SET_TOP(iter);
-            if (iter == NULL)
-                goto error;
+            PyObject *iter;
+            /* If we have a generator object on top -- keep it there,
+               it's already an iterator.
+
+               This is needed to allow use of 'async def' coroutines
+               in 'yield from' expression from generator-based coroutines
+               (decorated with types.coroutine()).
+
+               'yield from' is compiled to GET_ITER..YIELD_FROM combination,
+               but since coroutines raise TypeError in their 'tp_iter' we
+               need a way for them to "pass through" the GET_ITER.
+            */
+            if (!PyGen_CheckExact(iterable)) {
+                /* `iterable` is not a generator. */
+                iter = PyObject_GetIter(iterable);
+                Py_DECREF(iterable);
+                SET_TOP(iter);
+                if (iter == NULL)
+                    goto error;
+            }
             PREDICT(FOR_ITER);
             DISPATCH();
         }
@@ -2883,6 +3020,39 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
             DISPATCH();
         }
 
+        TARGET(BEFORE_ASYNC_WITH) {
+            _Py_IDENTIFIER(__aexit__);
+            _Py_IDENTIFIER(__aenter__);
+
+            PyObject *mgr = TOP();
+            PyObject *exit = special_lookup(mgr, &PyId___aexit__),
+                     *enter;
+            PyObject *res;
+            if (exit == NULL)
+                goto error;
+            SET_TOP(exit);
+            enter = special_lookup(mgr, &PyId___aenter__);
+            Py_DECREF(mgr);
+            if (enter == NULL)
+                goto error;
+            res = PyObject_CallFunctionObjArgs(enter, NULL);
+            Py_DECREF(enter);
+            if (res == NULL)
+                goto error;
+            PUSH(res);
+            DISPATCH();
+        }
+
+        TARGET(SETUP_ASYNC_WITH) {
+            PyObject *res = POP();
+            /* Setup the finally block before pushing the result
+               of __aenter__ on the stack. */
+            PyFrame_BlockSetup(f, SETUP_FINALLY, INSTR_OFFSET() + oparg,
+                               STACK_LEVEL());
+            PUSH(res);
+            DISPATCH();
+        }
+
         TARGET(SETUP_WITH) {
             _Py_IDENTIFIER(__exit__);
             _Py_IDENTIFIER(__enter__);
@@ -2909,7 +3079,7 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
             DISPATCH();
         }
 
-        TARGET(WITH_CLEANUP) {
+        TARGET(WITH_CLEANUP_START) {
             /* At the top of the stack are 1-6 values indicating
                how/why we entered the finally clause:
                - TOP = None
@@ -2937,7 +3107,6 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
 
             PyObject *exit_func;
             PyObject *exc = TOP(), *val = Py_None, *tb = Py_None, *res;
-            int err;
             if (exc == Py_None) {
                 (void)POP();
                 exit_func = TOP();
@@ -2987,10 +3156,23 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
             if (res == NULL)
                 goto error;
 
+            PUSH(exc);
+            PUSH(res);
+            PREDICT(WITH_CLEANUP_FINISH);
+            DISPATCH();
+        }
+
+        PREDICTED(WITH_CLEANUP_FINISH);
+        TARGET(WITH_CLEANUP_FINISH) {
+            PyObject *res = POP();
+            PyObject *exc = POP();
+            int err;
+
             if (exc != Py_None)
                 err = PyObject_IsTrue(res);
             else
                 err = 0;
+
             Py_DECREF(res);
 
             if (err < 0)
@@ -3751,6 +3933,9 @@ _PyEval_EvalCodeWithName(PyObject *_co, PyObject *globals, PyObject *locals,
     }
 
     if (co->co_flags & CO_GENERATOR) {
+        PyObject *gen;
+        PyObject *coroutine_wrapper;
+
         /* Don't need to keep the reference to f_back, it will be set
          * when the generator is resumed. */
         Py_CLEAR(f->f_back);
@@ -3759,7 +3944,19 @@ _PyEval_EvalCodeWithName(PyObject *_co, PyObject *globals, PyObject *locals,
 
         /* Create a new generator that owns the ready to run frame
          * and return that as the value. */
-        return PyGen_NewWithQualName(f, name, qualname);
+        gen = PyGen_NewWithQualName(f, name, qualname);
+        if (gen == NULL)
+            return NULL;
+
+        if (co->co_flags & (CO_COROUTINE | CO_ITERABLE_COROUTINE)) {
+            coroutine_wrapper = PyEval_GetCoroutineWrapper();
+            if (coroutine_wrapper != NULL) {
+                PyObject *wrapped =
+                            PyObject_CallFunction(coroutine_wrapper, "N", gen);
+                gen = wrapped;
+            }
+        }
+        return gen;
     }
 
     retval = PyEval_EvalFrameEx(f,0);
@@ -4205,6 +4402,24 @@ PyEval_SetTrace(Py_tracefunc func, PyObject *arg)
                            || (tstate->c_profilefunc != NULL));
 }
 
+void
+PyEval_SetCoroutineWrapper(PyObject *wrapper)
+{
+    PyThreadState *tstate = PyThreadState_GET();
+
+    Py_CLEAR(tstate->coroutine_wrapper);
+
+    Py_XINCREF(wrapper);
+    tstate->coroutine_wrapper = wrapper;
+}
+
+PyObject *
+PyEval_GetCoroutineWrapper()
+{
+    PyThreadState *tstate = PyThreadState_GET();
+    return tstate->coroutine_wrapper;
+}
+
 PyObject *
 PyEval_GetBuiltins(void)
 {