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

1 files changed, 273 insertions, 15 deletions

diff --git a/Python/compile.c b/Python/compile.c
index efd4807..7820f39 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -92,6 +92,7 @@ enum {
     COMPILER_SCOPE_MODULE,
     COMPILER_SCOPE_CLASS,
     COMPILER_SCOPE_FUNCTION,
+    COMPILER_SCOPE_ASYNC_FUNCTION,
     COMPILER_SCOPE_LAMBDA,
     COMPILER_SCOPE_COMPREHENSION,
 };
@@ -193,6 +194,8 @@ static int inplace_binop(struct compiler *, operator_ty);
 static int expr_constant(struct compiler *, expr_ty);
 
 static int compiler_with(struct compiler *, stmt_ty, int);
+static int compiler_async_with(struct compiler *, stmt_ty, int);
+static int compiler_async_for(struct compiler *, stmt_ty);
 static int compiler_call_helper(struct compiler *c, Py_ssize_t n,
                                 asdl_seq *args,
                                 asdl_seq *keywords);
@@ -673,7 +676,9 @@ compiler_set_qualname(struct compiler *c)
         parent = (struct compiler_unit *)PyCapsule_GetPointer(capsule, CAPSULE_NAME);
         assert(parent);
 
-        if (u->u_scope_type == COMPILER_SCOPE_FUNCTION || u->u_scope_type == COMPILER_SCOPE_CLASS) {
+        if (u->u_scope_type == COMPILER_SCOPE_FUNCTION
+            || u->u_scope_type == COMPILER_SCOPE_ASYNC_FUNCTION
+            || u->u_scope_type == COMPILER_SCOPE_CLASS) {
             assert(u->u_name);
             mangled = _Py_Mangle(parent->u_private, u->u_name);
             if (!mangled)
@@ -687,6 +692,7 @@ compiler_set_qualname(struct compiler *c)
 
         if (!force_global) {
             if (parent->u_scope_type == COMPILER_SCOPE_FUNCTION
+                || parent->u_scope_type == COMPILER_SCOPE_ASYNC_FUNCTION
                 || parent->u_scope_type == COMPILER_SCOPE_LAMBDA) {
                 dot_locals_str = _PyUnicode_FromId(&dot_locals);
                 if (dot_locals_str == NULL)
@@ -927,7 +933,9 @@ PyCompile_OpcodeStackEffect(int opcode, int oparg)
             return 0;
         case SETUP_WITH:
             return 7;
-        case WITH_CLEANUP:
+        case WITH_CLEANUP_START:
+            return 1;
+        case WITH_CLEANUP_FINISH:
             return -1; /* XXX Sometimes more */
         case RETURN_VALUE:
             return -1;
@@ -1048,6 +1056,16 @@ PyCompile_OpcodeStackEffect(int opcode, int oparg)
             return -1;
         case DELETE_DEREF:
             return 0;
+        case GET_AWAITABLE:
+            return 0;
+        case SETUP_ASYNC_WITH:
+            return 6;
+        case BEFORE_ASYNC_WITH:
+            return 1;
+        case GET_AITER:
+            return 0;
+        case GET_ANEXT:
+            return 1;
         default:
             return PY_INVALID_STACK_EFFECT;
     }
@@ -1642,19 +1660,43 @@ error:
 }
 
 static int
-compiler_function(struct compiler *c, stmt_ty s)
+compiler_function(struct compiler *c, stmt_ty s, int is_async)
 {
     PyCodeObject *co;
     PyObject *qualname, *first_const = Py_None;
-    arguments_ty args = s->v.FunctionDef.args;
-    expr_ty returns = s->v.FunctionDef.returns;
-    asdl_seq* decos = s->v.FunctionDef.decorator_list;
+    arguments_ty args;
+    expr_ty returns;
+    identifier name;
+    asdl_seq* decos;
+    asdl_seq *body;
     stmt_ty st;
     Py_ssize_t i, n, arglength;
     int docstring, kw_default_count = 0;
     int num_annotations;
+    int scope_type;
+
+
+    if (is_async) {
+        assert(s->kind == AsyncFunctionDef_kind);
+
+        args = s->v.AsyncFunctionDef.args;
+        returns = s->v.AsyncFunctionDef.returns;
+        decos = s->v.AsyncFunctionDef.decorator_list;
+        name = s->v.AsyncFunctionDef.name;
+        body = s->v.AsyncFunctionDef.body;
+
+        scope_type = COMPILER_SCOPE_ASYNC_FUNCTION;
+    } else {
+        assert(s->kind == FunctionDef_kind);
+
+        args = s->v.FunctionDef.args;
+        returns = s->v.FunctionDef.returns;
+        decos = s->v.FunctionDef.decorator_list;
+        name = s->v.FunctionDef.name;
+        body = s->v.FunctionDef.body;
 
-    assert(s->kind == FunctionDef_kind);
+        scope_type = COMPILER_SCOPE_FUNCTION;
+    }
 
     if (!compiler_decorators(c, decos))
         return 0;
@@ -1672,12 +1714,12 @@ compiler_function(struct compiler *c, stmt_ty s)
         return 0;
     assert((num_annotations & 0xFFFF) == num_annotations);
 
-    if (!compiler_enter_scope(c, s->v.FunctionDef.name,
-                              COMPILER_SCOPE_FUNCTION, (void *)s,
+    if (!compiler_enter_scope(c, name,
+                              scope_type, (void *)s,
                               s->lineno))
         return 0;
 
-    st = (stmt_ty)asdl_seq_GET(s->v.FunctionDef.body, 0);
+    st = (stmt_ty)asdl_seq_GET(body, 0);
     docstring = compiler_isdocstring(st);
     if (docstring && c->c_optimize < 2)
         first_const = st->v.Expr.value->v.Str.s;
@@ -1688,10 +1730,10 @@ compiler_function(struct compiler *c, stmt_ty s)
 
     c->u->u_argcount = asdl_seq_LEN(args->args);
     c->u->u_kwonlyargcount = asdl_seq_LEN(args->kwonlyargs);
-    n = asdl_seq_LEN(s->v.FunctionDef.body);
+    n = asdl_seq_LEN(body);
     /* if there was a docstring, we need to skip the first statement */
     for (i = docstring; i < n; i++) {
-        st = (stmt_ty)asdl_seq_GET(s->v.FunctionDef.body, i);
+        st = (stmt_ty)asdl_seq_GET(body, i);
         VISIT_IN_SCOPE(c, stmt, st);
     }
     co = assemble(c, 1);
@@ -1711,12 +1753,19 @@ compiler_function(struct compiler *c, stmt_ty s)
     Py_DECREF(qualname);
     Py_DECREF(co);
 
+    if (is_async) {
+        co->co_flags |= CO_COROUTINE;
+        /* An async function is always a generator, even
+           if there is no 'yield' expressions in it. */
+        co->co_flags |= CO_GENERATOR;
+    }
+
     /* decorators */
     for (i = 0; i < asdl_seq_LEN(decos); i++) {
         ADDOP_I(c, CALL_FUNCTION, 1);
     }
 
-    return compiler_nameop(c, s->v.FunctionDef.name, Store);
+    return compiler_nameop(c, name, Store);
 }
 
 static int
@@ -1989,6 +2038,92 @@ compiler_for(struct compiler *c, stmt_ty s)
     return 1;
 }
 
+
+static int
+compiler_async_for(struct compiler *c, stmt_ty s)
+{
+    static PyObject *stopiter_error = NULL;
+    basicblock *try, *except, *end, *after_try, *try_cleanup,
+               *after_loop, *after_loop_else;
+
+    if (stopiter_error == NULL) {
+        stopiter_error = PyUnicode_InternFromString("StopAsyncIteration");
+        if (stopiter_error == NULL)
+            return 0;
+    }
+
+    try = compiler_new_block(c);
+    except = compiler_new_block(c);
+    end = compiler_new_block(c);
+    after_try = compiler_new_block(c);
+    try_cleanup = compiler_new_block(c);
+    after_loop = compiler_new_block(c);
+    after_loop_else = compiler_new_block(c);
+
+    if (try == NULL || except == NULL || end == NULL
+            || after_try == NULL || try_cleanup == NULL)
+        return 0;
+
+    ADDOP_JREL(c, SETUP_LOOP, after_loop);
+    if (!compiler_push_fblock(c, LOOP, try))
+        return 0;
+
+    VISIT(c, expr, s->v.AsyncFor.iter);
+    ADDOP(c, GET_AITER);
+    ADDOP_O(c, LOAD_CONST, Py_None, consts);
+    ADDOP(c, YIELD_FROM);
+
+    compiler_use_next_block(c, try);
+
+
+    ADDOP_JREL(c, SETUP_EXCEPT, except);
+    if (!compiler_push_fblock(c, EXCEPT, try))
+        return 0;
+
+    ADDOP(c, GET_ANEXT);
+    ADDOP_O(c, LOAD_CONST, Py_None, consts);
+    ADDOP(c, YIELD_FROM);
+    VISIT(c, expr, s->v.AsyncFor.target);
+    ADDOP(c, POP_BLOCK);
+    compiler_pop_fblock(c, EXCEPT, try);
+    ADDOP_JREL(c, JUMP_FORWARD, after_try);
+
+
+    compiler_use_next_block(c, except);
+    ADDOP(c, DUP_TOP);
+    ADDOP_O(c, LOAD_GLOBAL, stopiter_error, names);
+    ADDOP_I(c, COMPARE_OP, PyCmp_EXC_MATCH);
+    ADDOP_JABS(c, POP_JUMP_IF_FALSE, try_cleanup);
+
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_EXCEPT); /* for SETUP_EXCEPT */
+    ADDOP(c, POP_BLOCK); /* for SETUP_LOOP */
+    ADDOP_JABS(c, JUMP_ABSOLUTE, after_loop_else);
+
+
+    compiler_use_next_block(c, try_cleanup);
+    ADDOP(c, END_FINALLY);
+
+    compiler_use_next_block(c, after_try);
+    VISIT_SEQ(c, stmt, s->v.AsyncFor.body);
+    ADDOP_JABS(c, JUMP_ABSOLUTE, try);
+
+    ADDOP(c, POP_BLOCK); /* for SETUP_LOOP */
+    compiler_pop_fblock(c, LOOP, try);
+
+    compiler_use_next_block(c, after_loop);
+    ADDOP_JABS(c, JUMP_ABSOLUTE, end);
+
+    compiler_use_next_block(c, after_loop_else);
+    VISIT_SEQ(c, stmt, s->v.For.orelse);
+
+    compiler_use_next_block(c, end);
+
+    return 1;
+}
+
 static int
 compiler_while(struct compiler *c, stmt_ty s)
 {
@@ -2515,7 +2650,7 @@ compiler_visit_stmt(struct compiler *c, stmt_ty s)
 
     switch (s->kind) {
     case FunctionDef_kind:
-        return compiler_function(c, s);
+        return compiler_function(c, s, 0);
     case ClassDef_kind:
         return compiler_class(c, s);
     case Return_kind:
@@ -2594,7 +2729,14 @@ compiler_visit_stmt(struct compiler *c, stmt_ty s)
         return compiler_continue(c);
     case With_kind:
         return compiler_with(c, s, 0);
+    case AsyncFunctionDef_kind:
+        return compiler_function(c, s, 1);
+    case AsyncWith_kind:
+        return compiler_async_with(c, s, 0);
+    case AsyncFor_kind:
+        return compiler_async_for(c, s);
     }
+
     return 1;
 }
 
@@ -3471,6 +3613,102 @@ expr_constant(struct compiler *c, expr_ty e)
     }
 }
 
+
+/*
+   Implements the async with statement.
+
+   The semantics outlined in that PEP are as follows:
+
+   async with EXPR as VAR:
+       BLOCK
+
+   It is implemented roughly as:
+
+   context = EXPR
+   exit = context.__aexit__  # not calling it
+   value = await context.__aenter__()
+   try:
+       VAR = value  # if VAR present in the syntax
+       BLOCK
+   finally:
+       if an exception was raised:
+       exc = copy of (exception, instance, traceback)
+       else:
+       exc = (None, None, None)
+       if not (await exit(*exc)):
+           raise
+ */
+static int
+compiler_async_with(struct compiler *c, stmt_ty s, int pos)
+{
+    basicblock *block, *finally;
+    withitem_ty item = asdl_seq_GET(s->v.AsyncWith.items, pos);
+
+    assert(s->kind == AsyncWith_kind);
+
+    block = compiler_new_block(c);
+    finally = compiler_new_block(c);
+    if (!block || !finally)
+        return 0;
+
+    /* Evaluate EXPR */
+    VISIT(c, expr, item->context_expr);
+
+    ADDOP(c, BEFORE_ASYNC_WITH);
+    ADDOP(c, GET_AWAITABLE);
+    ADDOP_O(c, LOAD_CONST, Py_None, consts);
+    ADDOP(c, YIELD_FROM);
+
+    ADDOP_JREL(c, SETUP_ASYNC_WITH, finally);
+
+    /* SETUP_ASYNC_WITH pushes a finally block. */
+    compiler_use_next_block(c, block);
+    if (!compiler_push_fblock(c, FINALLY_TRY, block)) {
+        return 0;
+    }
+
+    if (item->optional_vars) {
+        VISIT(c, expr, item->optional_vars);
+    }
+    else {
+    /* Discard result from context.__aenter__() */
+        ADDOP(c, POP_TOP);
+    }
+
+    pos++;
+    if (pos == asdl_seq_LEN(s->v.AsyncWith.items))
+        /* BLOCK code */
+        VISIT_SEQ(c, stmt, s->v.AsyncWith.body)
+    else if (!compiler_async_with(c, s, pos))
+            return 0;
+
+    /* End of try block; start the finally block */
+    ADDOP(c, POP_BLOCK);
+    compiler_pop_fblock(c, FINALLY_TRY, block);
+
+    ADDOP_O(c, LOAD_CONST, Py_None, consts);
+    compiler_use_next_block(c, finally);
+    if (!compiler_push_fblock(c, FINALLY_END, finally))
+        return 0;
+
+    /* Finally block starts; context.__exit__ is on the stack under
+       the exception or return information. Just issue our magic
+       opcode. */
+    ADDOP(c, WITH_CLEANUP_START);
+
+    ADDOP(c, GET_AWAITABLE);
+    ADDOP_O(c, LOAD_CONST, Py_None, consts);
+    ADDOP(c, YIELD_FROM);
+
+    ADDOP(c, WITH_CLEANUP_FINISH);
+
+    /* Finally block ends. */
+    ADDOP(c, END_FINALLY);
+    compiler_pop_fblock(c, FINALLY_END, finally);
+    return 1;
+}
+
+
 /*
    Implements the with statement from PEP 343.
 
@@ -3544,7 +3782,8 @@ compiler_with(struct compiler *c, stmt_ty s, int pos)
     /* Finally block starts; context.__exit__ is on the stack under
        the exception or return information. Just issue our magic
        opcode. */
-    ADDOP(c, WITH_CLEANUP);
+    ADDOP(c, WITH_CLEANUP_START);
+    ADDOP(c, WITH_CLEANUP_FINISH);
 
     /* Finally block ends. */
     ADDOP(c, END_FINALLY);
@@ -3595,6 +3834,8 @@ 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);
         }
@@ -3606,11 +3847,28 @@ compiler_visit_expr(struct compiler *c, expr_ty e)
     case YieldFrom_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 from' inside async function");
+
         VISIT(c, expr, e->v.YieldFrom.value);
         ADDOP(c, GET_ITER);
         ADDOP_O(c, LOAD_CONST, Py_None, consts);
         ADDOP(c, YIELD_FROM);
         break;
+    case Await_kind:
+        if (c->u->u_ste->ste_type != FunctionBlock)
+            return compiler_error(c, "'await' outside function");
+
+        /* this check won't be triggered while we have AWAIT token */
+        if (c->u->u_scope_type != COMPILER_SCOPE_ASYNC_FUNCTION)
+            return compiler_error(c, "'await' outside async function");
+
+        VISIT(c, expr, e->v.Await.value);
+        ADDOP(c, GET_AWAITABLE);
+        ADDOP_O(c, LOAD_CONST, Py_None, consts);
+        ADDOP(c, YIELD_FROM);
+        break;
     case Compare_kind:
         return compiler_compare(c, e);
     case Call_kind: