1 files changed, 39 insertions, 40 deletions

diff --git a/Python/ceval.c b/Python/ceval.c
index e127a731fd..12df0fe3fc 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -1191,7 +1191,7 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
     f->f_stacktop = NULL;       /* remains NULL unless yield suspends frame */
     f->f_executing = 1;
 
-    if (co->co_flags & CO_GENERATOR) {
+    if (co->co_flags & (CO_GENERATOR | CO_COROUTINE)) {
         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
@@ -1955,7 +1955,7 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
                 goto error;
             }
 
-            awaitable = _PyGen_GetAwaitableIter(iter);
+            awaitable = _PyCoro_GetAwaitableIter(iter);
             if (awaitable == NULL) {
                 SET_TOP(NULL);
                 PyErr_Format(
@@ -1998,7 +1998,7 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
                 goto error;
             }
 
-            awaitable = _PyGen_GetAwaitableIter(next_iter);
+            awaitable = _PyCoro_GetAwaitableIter(next_iter);
             if (awaitable == NULL) {
                 PyErr_Format(
                     PyExc_TypeError,
@@ -2017,7 +2017,7 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
 
         TARGET(GET_AWAITABLE) {
             PyObject *iterable = TOP();
-            PyObject *iter = _PyGen_GetAwaitableIter(iterable);
+            PyObject *iter = _PyCoro_GetAwaitableIter(iterable);
 
             Py_DECREF(iterable);
 
@@ -2034,25 +2034,7 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
             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;
-                }
+            if (PyGen_CheckExact(reciever) || PyCoro_CheckExact(reciever)) {
                 retval = _PyGen_Send((PyGenObject *)reciever, v);
             } else {
                 _Py_IDENTIFIER(send);
@@ -2929,19 +2911,33 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
         TARGET(GET_ITER) {
             /* before: [obj]; after [getiter(obj)] */
             PyObject *iterable = TOP();
-            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()).
+            PyObject *iter = PyObject_GetIter(iterable);
+            Py_DECREF(iterable);
+            SET_TOP(iter);
+            if (iter == NULL)
+                goto error;
+            PREDICT(FOR_ITER);
+            DISPATCH();
+        }
 
-               '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)) {
+        TARGET(GET_YIELD_FROM_ITER) {
+            /* before: [obj]; after [getiter(obj)] */
+            PyObject *iterable = TOP();
+            PyObject *iter;
+            if (PyCoro_CheckExact(iterable)) {
+                /* `iterable` is a coroutine */
+                if (!(co->co_flags & (CO_COROUTINE | CO_ITERABLE_COROUTINE))) {
+                    /* and it is used in a 'yield from' expression of a
+                       regular generator. */
+                    Py_DECREF(iterable);
+                    SET_TOP(NULL);
+                    PyErr_SetString(PyExc_TypeError,
+                                    "cannot 'yield from' a coroutine object "
+                                    "in a non-coroutine generator");
+                    goto error;
+                }
+            }
+            else if (!PyGen_CheckExact(iterable)) {
                 /* `iterable` is not a generator. */
                 iter = PyObject_GetIter(iterable);
                 Py_DECREF(iterable);
@@ -2949,7 +2945,6 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
                 if (iter == NULL)
                     goto error;
             }
-            PREDICT(FOR_ITER);
             DISPATCH();
         }
 
@@ -3517,7 +3512,7 @@ fast_block_end:
     assert((retval != NULL) ^ (PyErr_Occurred() != NULL));
 
 fast_yield:
-    if (co->co_flags & CO_GENERATOR) {
+    if (co->co_flags & (CO_GENERATOR | CO_COROUTINE)) {
 
         /* The purpose of this block is to put aside the generator's exception
            state and restore that of the calling frame. If the current
@@ -3919,10 +3914,10 @@ _PyEval_EvalCodeWithName(PyObject *_co, PyObject *globals, PyObject *locals,
         freevars[PyTuple_GET_SIZE(co->co_cellvars) + i] = o;
     }
 
-    if (co->co_flags & CO_GENERATOR) {
+    if (co->co_flags & (CO_GENERATOR | CO_COROUTINE)) {
         PyObject *gen;
         PyObject *coro_wrapper = tstate->coroutine_wrapper;
-        int is_coro = co->co_flags & (CO_COROUTINE | CO_ITERABLE_COROUTINE);
+        int is_coro = co->co_flags & CO_COROUTINE;
 
         if (is_coro && tstate->in_coroutine_wrapper) {
             assert(coro_wrapper != NULL);
@@ -3942,7 +3937,11 @@ _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. */
-        gen = PyGen_NewWithQualName(f, name, qualname);
+        if (is_coro) {
+            gen = PyCoro_New(f, name, qualname);
+        } else {
+            gen = PyGen_NewWithQualName(f, name, qualname);
+        }
         if (gen == NULL)
             return NULL;