PEP 0492 -- Coroutines with async and await syntax. Issue #24017.

6 files changed, 198 insertions, 17 deletions

diff --git a/Objects/exceptions.c b/Objects/exceptions.c
index e09d384..d494995 100644
--- a/Objects/exceptions.c
+++ b/Objects/exceptions.c
@@ -473,6 +473,13 @@ SimpleExtendsException(PyExc_Exception, TypeError,
 
 
 /*
+ *    StopAsyncIteration extends Exception
+ */
+SimpleExtendsException(PyExc_Exception, StopAsyncIteration,
+                       "Signal the end from iterator.__anext__().");
+
+
+/*
  *    StopIteration extends Exception
  */
 
@@ -2468,6 +2475,7 @@ _PyExc_Init(PyObject *bltinmod)
     PRE_INIT(BaseException)
     PRE_INIT(Exception)
     PRE_INIT(TypeError)
+    PRE_INIT(StopAsyncIteration)
     PRE_INIT(StopIteration)
     PRE_INIT(GeneratorExit)
     PRE_INIT(SystemExit)
@@ -2538,6 +2546,7 @@ _PyExc_Init(PyObject *bltinmod)
     POST_INIT(BaseException)
     POST_INIT(Exception)
     POST_INIT(TypeError)
+    POST_INIT(StopAsyncIteration)
     POST_INIT(StopIteration)
     POST_INIT(GeneratorExit)
     POST_INIT(SystemExit)
diff --git a/Objects/frameobject.c b/Objects/frameobject.c
index 55ee563..172f2cb 100644
--- a/Objects/frameobject.c
+++ b/Objects/frameobject.c
@@ -196,6 +196,7 @@ frame_setlineno(PyFrameObject *f, PyObject* p_new_lineno)
         case SETUP_EXCEPT:
         case SETUP_FINALLY:
         case SETUP_WITH:
+        case SETUP_ASYNC_WITH:
             blockstack[blockstack_top++] = addr;
             in_finally[blockstack_top-1] = 0;
             break;
@@ -203,7 +204,8 @@ frame_setlineno(PyFrameObject *f, PyObject* p_new_lineno)
         case POP_BLOCK:
             assert(blockstack_top > 0);
             setup_op = code[blockstack[blockstack_top-1]];
-            if (setup_op == SETUP_FINALLY || setup_op == SETUP_WITH) {
+            if (setup_op == SETUP_FINALLY || setup_op == SETUP_WITH
+                                    || setup_op == SETUP_ASYNC_WITH) {
                 in_finally[blockstack_top-1] = 1;
             }
             else {
@@ -218,7 +220,8 @@ frame_setlineno(PyFrameObject *f, PyObject* p_new_lineno)
              * be seeing such an END_FINALLY.) */
             if (blockstack_top > 0) {
                 setup_op = code[blockstack[blockstack_top-1]];
-                if (setup_op == SETUP_FINALLY || setup_op == SETUP_WITH) {
+                if (setup_op == SETUP_FINALLY || setup_op == SETUP_WITH
+                                    || setup_op == SETUP_ASYNC_WITH) {
                     blockstack_top--;
                 }
             }
@@ -281,6 +284,7 @@ frame_setlineno(PyFrameObject *f, PyObject* p_new_lineno)
         case SETUP_EXCEPT:
         case SETUP_FINALLY:
         case SETUP_WITH:
+        case SETUP_ASYNC_WITH:
             delta_iblock++;
             break;
 
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;
+}
diff --git a/Objects/typeobject.c b/Objects/typeobject.c
index 4b99287..9522ac5 100644
--- a/Objects/typeobject.c
+++ b/Objects/typeobject.c
@@ -2506,6 +2506,7 @@ type_new(PyTypeObject *metatype, PyObject *args, PyObject *kwds)
         type->tp_flags |= Py_TPFLAGS_HAVE_GC;
 
     /* Initialize essential fields */
+    type->tp_as_async = &et->as_async;
     type->tp_as_number = &et->as_number;
     type->tp_as_sequence = &et->as_sequence;
     type->tp_as_mapping = &et->as_mapping;
@@ -2751,6 +2752,7 @@ PyType_FromSpecWithBases(PyType_Spec *spec, PyObject *bases)
     }
 
     /* Initialize essential fields */
+    type->tp_as_async = &res->as_async;
     type->tp_as_number = &res->as_number;
     type->tp_as_sequence = &res->as_sequence;
     type->tp_as_mapping = &res->as_mapping;
@@ -4566,6 +4568,7 @@ inherit_slots(PyTypeObject *type, PyTypeObject *base)
 #define COPYSLOT(SLOT) \
     if (!type->SLOT && SLOTDEFINED(SLOT)) type->SLOT = base->SLOT
 
+#define COPYASYNC(SLOT) COPYSLOT(tp_as_async->SLOT)
 #define COPYNUM(SLOT) COPYSLOT(tp_as_number->SLOT)
 #define COPYSEQ(SLOT) COPYSLOT(tp_as_sequence->SLOT)
 #define COPYMAP(SLOT) COPYSLOT(tp_as_mapping->SLOT)
@@ -4615,6 +4618,15 @@ inherit_slots(PyTypeObject *type, PyTypeObject *base)
         COPYNUM(nb_inplace_matrix_multiply);
     }
 
+    if (type->tp_as_async != NULL && base->tp_as_async != NULL) {
+        basebase = base->tp_base;
+        if (basebase->tp_as_async == NULL)
+            basebase = NULL;
+        COPYASYNC(am_await);
+        COPYASYNC(am_aiter);
+        COPYASYNC(am_anext);
+    }
+
     if (type->tp_as_sequence != NULL && base->tp_as_sequence != NULL) {
         basebase = base->tp_base;
         if (basebase->tp_as_sequence == NULL)
@@ -4884,6 +4896,8 @@ PyType_Ready(PyTypeObject *type)
     /* Some more special stuff */
     base = type->tp_base;
     if (base != NULL) {
+        if (type->tp_as_async == NULL)
+            type->tp_as_async = base->tp_as_async;
         if (type->tp_as_number == NULL)
             type->tp_as_number = base->tp_as_number;
         if (type->tp_as_sequence == NULL)
@@ -4904,16 +4918,6 @@ PyType_Ready(PyTypeObject *type)
             goto error;
     }
 
-    /* Warn for a type that implements tp_compare (now known as
-       tp_reserved) but not tp_richcompare. */
-    if (type->tp_reserved && !type->tp_richcompare) {
-        PyErr_Format(PyExc_TypeError,
-            "Type %.100s defines tp_reserved (formerly tp_compare) "
-            "but not tp_richcompare. Comparisons may not behave as intended.",
-            type->tp_name);
-        goto error;
-    }
-
     /* All done -- set the ready flag */
     assert(type->tp_dict != NULL);
     type->tp_flags =
@@ -6265,6 +6269,59 @@ slot_tp_finalize(PyObject *self)
     PyErr_Restore(error_type, error_value, error_traceback);
 }
 
+static PyObject *
+slot_am_await(PyObject *self)
+{
+    PyObject *func, *res;
+    _Py_IDENTIFIER(__await__);
+
+    func = lookup_method(self, &PyId___await__);
+    if (func != NULL) {
+        res = PyEval_CallObject(func, NULL);
+        Py_DECREF(func);
+        return res;
+    }
+    PyErr_Format(PyExc_AttributeError,
+                 "object %.50s does not have __await__ method",
+                 Py_TYPE(self)->tp_name);
+    return NULL;
+}
+
+static PyObject *
+slot_am_aiter(PyObject *self)
+{
+    PyObject *func, *res;
+    _Py_IDENTIFIER(__aiter__);
+
+    func = lookup_method(self, &PyId___aiter__);
+    if (func != NULL) {
+        res = PyEval_CallObject(func, NULL);
+        Py_DECREF(func);
+        return res;
+    }
+    PyErr_Format(PyExc_AttributeError,
+                 "object %.50s does not have __aiter__ method",
+                 Py_TYPE(self)->tp_name);
+    return NULL;
+}
+
+static PyObject *
+slot_am_anext(PyObject *self)
+{
+    PyObject *func, *res;
+    _Py_IDENTIFIER(__anext__);
+
+    func = lookup_method(self, &PyId___anext__);
+    if (func != NULL) {
+        res = PyEval_CallObject(func, NULL);
+        Py_DECREF(func);
+        return res;
+    }
+    PyErr_Format(PyExc_AttributeError,
+                 "object %.50s does not have __anext__ method",
+                 Py_TYPE(self)->tp_name);
+    return NULL;
+}
 
 /*
 Table mapping __foo__ names to tp_foo offsets and slot_tp_foo wrapper functions.
@@ -6281,6 +6338,7 @@ typedef struct wrapperbase slotdef;
 
 #undef TPSLOT
 #undef FLSLOT
+#undef AMSLOT
 #undef ETSLOT
 #undef SQSLOT
 #undef MPSLOT
@@ -6299,6 +6357,8 @@ typedef struct wrapperbase slotdef;
 #define ETSLOT(NAME, SLOT, FUNCTION, WRAPPER, DOC) \
     {NAME, offsetof(PyHeapTypeObject, SLOT), (void *)(FUNCTION), WRAPPER, \
      PyDoc_STR(DOC)}
+#define AMSLOT(NAME, SLOT, FUNCTION, WRAPPER, DOC) \
+    ETSLOT(NAME, as_async.SLOT, FUNCTION, WRAPPER, DOC)
 #define SQSLOT(NAME, SLOT, FUNCTION, WRAPPER, DOC) \
     ETSLOT(NAME, as_sequence.SLOT, FUNCTION, WRAPPER, DOC)
 #define MPSLOT(NAME, SLOT, FUNCTION, WRAPPER, DOC) \
@@ -6378,6 +6438,13 @@ static slotdef slotdefs[] = {
            "Create and return new object.  See help(type) for accurate signature."),
     TPSLOT("__del__", tp_finalize, slot_tp_finalize, (wrapperfunc)wrap_del, ""),
 
+    AMSLOT("__await__", am_await, slot_am_await, wrap_unaryfunc,
+           "__await__($self, /)\n--\n\nReturn an iterator to be used in await expression."),
+    AMSLOT("__aiter__", am_aiter, slot_am_aiter, wrap_unaryfunc,
+           "__aiter__($self, /)\n--\n\nReturn an awaitable, that resolves in asynchronous iterator."),
+    AMSLOT("__anext__", am_anext, slot_am_anext, wrap_unaryfunc,
+           "__anext__($self, /)\n--\n\nReturn a value or raise StopAsyncIteration."),
+
     BINSLOT("__add__", nb_add, slot_nb_add,
            "+"),
     RBINSLOT("__radd__", nb_add, slot_nb_add,
@@ -6530,6 +6597,10 @@ slotptr(PyTypeObject *type, int ioffset)
         ptr = (char *)type->tp_as_number;
         offset -= offsetof(PyHeapTypeObject, as_number);
     }
+    else if ((size_t)offset >= offsetof(PyHeapTypeObject, as_async)) {
+        ptr = (char *)type->tp_as_async;
+        offset -= offsetof(PyHeapTypeObject, as_async);
+    }
     else {
         ptr = (char *)type;
     }
diff --git a/Objects/typeslots.inc b/Objects/typeslots.inc
index 2ed99d8..bcbe0af 100644
--- a/Objects/typeslots.inc
+++ b/Objects/typeslots.inc
@@ -75,3 +75,6 @@ offsetof(PyHeapTypeObject, ht_type.tp_getset),
 offsetof(PyHeapTypeObject, ht_type.tp_free),
 offsetof(PyHeapTypeObject, as_number.nb_matrix_multiply),
 offsetof(PyHeapTypeObject, as_number.nb_inplace_matrix_multiply),
+offsetof(PyHeapTypeObject, as_async.am_await),
+offsetof(PyHeapTypeObject, as_async.am_aiter),
+offsetof(PyHeapTypeObject, as_async.am_anext),
diff --git a/Objects/typeslots.py b/Objects/typeslots.py
index b24c7f4..ba37c40 100755
--- a/Objects/typeslots.py
+++ b/Objects/typeslots.py
@@ -12,6 +12,8 @@ for line in sys.stdin:
     member = m.group(1)
     if member.startswith("tp_"):
         member = "ht_type."+member
+    elif member.startswith("am_"):
+        member = "as_async."+member
     elif member.startswith("nb_"):
         member = "as_number."+member
     elif member.startswith("mp_"):