1 files changed, 97 insertions, 5 deletions

diff --git a/Objects/genobject.c b/Objects/genobject.c
index 82b629c..2a5a0d9 100644
--- a/Objects/genobject.c
+++ b/Objects/genobject.c
@@ -24,6 +24,19 @@ _PyGen_Finalize(PyObject *self)
     PyObject *res;
     PyObject *error_type, *error_value, *error_traceback;
 
+    /* If `gen` is a coroutine, and if it was never awaited on,
+       issue a RuntimeWarning. */
+    if (gen->gi_code != NULL
+            && ((PyCodeObject *)gen->gi_code)->co_flags & (CO_COROUTINE
+                                                       | CO_ITERABLE_COROUTINE)
+            && gen->gi_frame != NULL
+            && gen->gi_frame->f_lasti == -1
+            && !PyErr_Occurred()
+            && PyErr_WarnFormat(PyExc_RuntimeWarning, 1,
+                                "coroutine '%.50S' was never awaited",
+                                gen->gi_qualname))
+        return;
+
     if (gen->gi_frame == NULL || gen->gi_frame->f_stacktop == NULL)
         /* Generator isn't paused, so no need to close */
         return;
@@ -135,7 +148,7 @@ gen_send_ex(PyGenObject *gen, PyObject *arg, int exc)
          * a leaking StopIteration into RuntimeError (with its cause
          * set appropriately). */
         if ((((PyCodeObject *)gen->gi_code)->co_flags &
-                                                CO_FUTURE_GENERATOR_STOP)
+              (CO_FUTURE_GENERATOR_STOP | CO_COROUTINE | CO_ITERABLE_COROUTINE))
             && PyErr_ExceptionMatches(PyExc_StopIteration))
         {
             PyObject *exc, *val, *val2, *tb;
@@ -402,6 +415,12 @@ failed_throw:
 static PyObject *
 gen_iternext(PyGenObject *gen)
 {
+    if (((PyCodeObject*)gen->gi_code)->co_flags & CO_COROUTINE) {
+        PyErr_SetString(PyExc_TypeError,
+                        "coroutine-objects do not support iteration");
+        return NULL;
+    }
+
     return gen_send_ex(gen, NULL, 0);
 }
 
@@ -459,8 +478,14 @@ _PyGen_FetchStopIterationValue(PyObject **pvalue) {
 static PyObject *
 gen_repr(PyGenObject *gen)
 {
-    return PyUnicode_FromFormat("<generator object %S at %p>",
-                                gen->gi_qualname, gen);
+    if (PyGen_CheckCoroutineExact(gen)) {
+        return PyUnicode_FromFormat("<coroutine object %S at %p>",
+                                    gen->gi_qualname, gen);
+    }
+    else {
+        return PyUnicode_FromFormat("<generator object %S at %p>",
+                                    gen->gi_qualname, gen);
+    }
 }
 
 static PyObject *
@@ -496,6 +521,19 @@ gen_get_qualname(PyGenObject *op)
     return op->gi_qualname;
 }
 
+static PyObject *
+gen_get_iter(PyGenObject *gen)
+{
+    if (((PyCodeObject*)gen->gi_code)->co_flags & CO_COROUTINE) {
+        PyErr_SetString(PyExc_TypeError,
+                        "coroutine-objects do not support iteration");
+        return NULL;
+    }
+
+    Py_INCREF(gen);
+    return gen;
+}
+
 static int
 gen_set_qualname(PyGenObject *op, PyObject *value)
 {
@@ -547,7 +585,7 @@ PyTypeObject PyGen_Type = {
     0,                                          /* tp_print */
     0,                                          /* tp_getattr */
     0,                                          /* tp_setattr */
-    0,                                          /* tp_reserved */
+    0,                                          /* tp_as_async */
     (reprfunc)gen_repr,                         /* tp_repr */
     0,                                          /* tp_as_number */
     0,                                          /* tp_as_sequence */
@@ -565,7 +603,7 @@ PyTypeObject PyGen_Type = {
     0,                                          /* tp_clear */
     0,                                          /* tp_richcompare */
     offsetof(PyGenObject, gi_weakreflist),      /* tp_weaklistoffset */
-    PyObject_SelfIter,                          /* tp_iter */
+    (getiterfunc)gen_get_iter,                  /* tp_iter */
     (iternextfunc)gen_iternext,                 /* tp_iternext */
     gen_methods,                                /* tp_methods */
     gen_memberlist,                             /* tp_members */
@@ -642,3 +680,57 @@ PyGen_NeedsFinalizing(PyGenObject *gen)
     /* No blocks except loops, it's safe to skip finalization. */
     return 0;
 }
+
+/*
+ *   This helper function returns an awaitable for `o`:
+ *     - `o` if `o` is a coroutine-object;
+ *     - `type(o)->tp_as_async->am_await(o)`
+ *
+ *   Raises a TypeError if it's not possible to return
+ *   an awaitable and returns NULL.
+ */
+PyObject *
+_PyGen_GetAwaitableIter(PyObject *o)
+{
+    getawaitablefunc getter = NULL;
+    PyTypeObject *ot;
+
+    if (PyGen_CheckCoroutineExact(o)) {
+        /* Fast path. It's a central function for 'await'. */
+        Py_INCREF(o);
+        return o;
+    }
+
+    ot = Py_TYPE(o);
+    if (ot->tp_as_async != NULL) {
+        getter = ot->tp_as_async->am_await;
+    }
+    if (getter != NULL) {
+        PyObject *res = (*getter)(o);
+        if (res != NULL) {
+            if (!PyIter_Check(res)) {
+                PyErr_Format(PyExc_TypeError,
+                             "__await__() returned non-iterator "
+                             "of type '%.100s'",
+                             Py_TYPE(res)->tp_name);
+                Py_CLEAR(res);
+            }
+            else {
+                if (PyGen_CheckCoroutineExact(res)) {
+                    /* __await__ must return an *iterator*, not
+                       a coroutine or another awaitable (see PEP 492) */
+                    PyErr_SetString(PyExc_TypeError,
+                                    "__await__() returned a coroutine");
+                    Py_CLEAR(res);
+                }
+            }
+        }
+        return res;
+    }
+
+    PyErr_Format(PyExc_TypeError,
+                 "object %.100s can't be used in 'await' expression",
+                 ot->tp_name);
+
+    return NULL;
+}