micropython · dpgeorge · Nov 27, 2020
diff --git a/ports/unix/coverage.c b/ports/unix/coverage.c
@@ -552,7 +552,9 @@ STATIC mp_obj_t extra_coverage(void) {
         code_state->exc_sp_idx = 0;
         code_state->old_globals = NULL;
         mp_vm_return_kind_t ret = mp_execute_bytecode(code_state, MP_OBJ_NULL);
-        mp_printf(&mp_plat_print, "%d %d\n", ret, mp_obj_get_type(code_state->state[0]) == &mp_type_NotImplementedError);
+        mp_printf(&mp_plat_print, "%d %d\n",
+            ret == MP_VM_RETURN_EXCEPTION,
+            mp_obj_get_type(code_state->state[0]) == &mp_type_NotImplementedError);
     }
 
     // scheduler

diff --git a/py/compile.c b/py/compile.c
@@ -938,7 +938,7 @@ STATIC void compile_decorated(compiler_t *comp, mp_parse_node_struct_t *pns) {
         mp_parse_node_struct_t *pns0 = (mp_parse_node_struct_t *)pns_body->nodes[0];
         body_name = compile_funcdef_helper(comp, pns0, emit_options);
         scope_t *fscope = (scope_t *)pns0->nodes[4];
-        fscope->scope_flags |= MP_SCOPE_FLAG_GENERATOR;
+        fscope->scope_flags |= MP_SCOPE_FLAG_GENERATOR | MP_SCOPE_FLAG_ASYNCDEF;
     #endif
     } else {
         assert(MP_PARSE_NODE_STRUCT_KIND(pns_body) == PN_classdef); // should be
@@ -1041,6 +1041,12 @@ STATIC void compile_return_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
         return;
     }
     #endif
+    #if MICROPY_PY_ASYNC_AWAIT
+    if ((comp->scope_cur->scope_flags & MP_SCOPE_FLAG_ASYNCGENERATOR)
+        && !MP_PARSE_NODE_IS_NULL(pns->nodes[0])) {
+        compile_syntax_error(comp, (mp_parse_node_t)pns, MP_ERROR_TEXT("'return' with value in async generator"));
+    }
+    #endif
     if (MP_PARSE_NODE_IS_NULL(pns->nodes[0])) {
         // no argument to 'return', so return None
         EMIT_ARG(load_const_tok, MP_TOKEN_KW_NONE);
@@ -1765,6 +1771,7 @@ STATIC void compile_with_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
 }
 
 STATIC void compile_yield_from(compiler_t *comp) {
+    comp->scope_cur->scope_flags |= MP_SCOPE_FLAG_GENERATOR;
     EMIT_ARG(get_iter, false);
     EMIT_ARG(load_const_tok, MP_TOKEN_KW_NONE);
     EMIT_ARG(yield, MP_EMIT_YIELD_FROM);
@@ -1959,7 +1966,7 @@ STATIC void compile_async_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
         // async def
         compile_funcdef(comp, pns0);
         scope_t *fscope = (scope_t *)pns0->nodes[4];
-        fscope->scope_flags |= MP_SCOPE_FLAG_GENERATOR;
+        fscope->scope_flags |= MP_SCOPE_FLAG_GENERATOR | MP_SCOPE_FLAG_ASYNCDEF;
     } else {
         // async for/with; first verify the scope is a generator
         int scope_flags = comp->scope_cur->scope_flags;
@@ -2711,6 +2718,12 @@ STATIC void compile_yield_expr(compiler_t *comp, mp_parse_node_struct_t *pns) {
         compile_syntax_error(comp, (mp_parse_node_t)pns, MP_ERROR_TEXT("'yield' outside function"));
         return;
     }
+    comp->scope_cur->scope_flags |= MP_SCOPE_FLAG_GENERATOR;
+    #if MICROPY_PY_ASYNC_AWAIT
+    if (comp->scope_cur->scope_flags & MP_SCOPE_FLAG_ASYNCDEF) {
+        comp->scope_cur->scope_flags |= MP_SCOPE_FLAG_ASYNCGENERATOR;
+    }
+    #endif
     if (MP_PARSE_NODE_IS_NULL(pns->nodes[0])) {
         EMIT_ARG(load_const_tok, MP_TOKEN_KW_NONE);
         EMIT_ARG(yield, MP_EMIT_YIELD_VALUE);
@@ -2719,6 +2732,11 @@ STATIC void compile_yield_expr(compiler_t *comp, mp_parse_node_struct_t *pns) {
         pns = (mp_parse_node_struct_t *)pns->nodes[0];
         compile_node(comp, pns->nodes[0]);
         compile_yield_from(comp);
+        #if MICROPY_PY_ASYNC_AWAIT
+        if (comp->scope_cur->scope_flags & MP_SCOPE_FLAG_ASYNCDEF) {
+            compile_syntax_error(comp, (mp_parse_node_t)pns, MP_ERROR_TEXT("'yield from' in async function"));
+        }
+        #endif
     } else {
         compile_node(comp, pns->nodes[0]);
         EMIT_ARG(yield, MP_EMIT_YIELD_VALUE);
@@ -2920,6 +2938,7 @@ STATIC void compile_scope_comp_iter(compiler_t *comp, mp_parse_node_struct_t *pn
         // no more nested if/for; compile inner expression
         compile_node(comp, pn_inner_expr);
         if (comp->scope_cur->kind == SCOPE_GEN_EXPR) {
+            comp->scope_cur->scope_flags |= MP_SCOPE_FLAG_GENERATOR;
             EMIT_ARG(yield, MP_EMIT_YIELD_VALUE);
             reserve_labels_for_native(comp, 1);
             EMIT(pop_top);

diff --git a/py/emitbc.c b/py/emitbc.c
@@ -800,7 +800,6 @@ void mp_emit_bc_raise_varargs(emit_t *emit, mp_uint_t n_args) {
 void mp_emit_bc_yield(emit_t *emit, int kind) {
     MP_STATIC_ASSERT(MP_BC_YIELD_VALUE + 1 == MP_BC_YIELD_FROM);
     emit_write_bytecode_byte(emit, -kind, MP_BC_YIELD_VALUE + kind);
-    emit->scope->scope_flags |= MP_SCOPE_FLAG_GENERATOR;
 }
 
 void mp_emit_bc_start_except_handler(emit_t *emit) {

diff --git a/py/emitnative.c b/py/emitnative.c
@@ -2831,7 +2831,6 @@ STATIC void emit_native_yield(emit_t *emit, int kind) {
     if (emit->do_viper_types) {
         mp_raise_NotImplementedError(MP_ERROR_TEXT("native yield"));
     }
-    emit->scope->scope_flags |= MP_SCOPE_FLAG_GENERATOR;
 
     need_stack_settled(emit);
 
@@ -2853,7 +2852,13 @@ STATIC void emit_native_yield(emit_t *emit, int kind) {
     emit_native_mov_state_reg(emit, OFFSETOF_CODE_STATE_SP, REG_TEMP0);
 
     // Put return type in return value slot
-    ASM_MOV_REG_IMM(emit->as, REG_TEMP0, MP_VM_RETURN_YIELD);
+    mp_vm_return_kind_t ret_kind = MP_VM_RETURN_YIELD;
+    #if MICROPY_PY_ASYNC_AWAIT
+    if ((emit->scope->scope_flags & MP_SCOPE_FLAG_ASYNCGENERATOR) && kind == MP_EMIT_YIELD_FROM) {
+        ret_kind = MP_VM_RETURN_YIELD_FROM;
+    }
+    #endif
+    ASM_MOV_REG_IMM(emit->as, REG_TEMP0, ret_kind);
     ASM_MOV_LOCAL_REG(emit->as, LOCAL_IDX_RET_VAL(emit), REG_TEMP0);
 
     // Save re-entry PC

diff --git a/py/obj.h b/py/obj.h
@@ -813,6 +813,7 @@ extern const mp_obj_type_t mp_type_super;
 extern const mp_obj_type_t mp_type_gen_wrap;
 extern const mp_obj_type_t mp_type_native_gen_wrap;
 extern const mp_obj_type_t mp_type_gen_instance;
+extern const mp_obj_type_t mp_type_agen_instance;
 extern const mp_obj_type_t mp_type_fun_builtin_0;
 extern const mp_obj_type_t mp_type_fun_builtin_1;
 extern const mp_obj_type_t mp_type_fun_builtin_2;

diff --git a/py/objgenerator.c b/py/objgenerator.c
@@ -3,7 +3,7 @@
  *
  * The MIT License (MIT)
  *
- * Copyright (c) 2013-2019 Damien P. George
+ * Copyright (c) 2013-2020 Damien P. George
  * Copyright (c) 2014-2017 Paul Sokolovsky
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
@@ -29,6 +29,7 @@
 #include <assert.h>
 
 #include "py/runtime.h"
+#include "py/bc0.h"
 #include "py/bc.h"
 #include "py/objstr.h"
 #include "py/objgenerator.h"
@@ -63,6 +64,12 @@ STATIC mp_obj_t gen_wrap_call(mp_obj_t self_in, size_t n_args, size_t n_kw, cons
         n_state * sizeof(mp_obj_t) + n_exc_stack * sizeof(mp_exc_stack_t),
         &mp_type_gen_instance);
 
+    #if MICROPY_PY_ASYNC_AWAIT
+    if (scope_flags & MP_SCOPE_FLAG_ASYNCGENERATOR) {
+        o->base.type = &mp_type_agen_instance;
+    }
+    #endif
+
     o->pend_exc = mp_const_none;
     o->code_state.fun_bc = self_fun;
     o->code_state.n_state = n_state;
@@ -116,6 +123,11 @@ STATIC mp_obj_t native_gen_wrap_call(mp_obj_t self_in, size_t n_args, size_t n_k
 
     // Allocate the generator object, with room for local stack (exception stack not needed).
     mp_obj_gen_instance_native_t *o = mp_obj_malloc_var(mp_obj_gen_instance_native_t, byte, n_state * sizeof(mp_obj_t), &mp_type_gen_instance);
+    #if MICROPY_PY_ASYNC_AWAIT
+    if (scope_flags & MP_SCOPE_FLAG_ASYNCGENERATOR) {
+        o->base.type = &mp_type_agen_instance;
+    }
+    #endif
 
     // Parse the input arguments and set up the code state
     o->pend_exc = mp_const_none;
@@ -167,6 +179,13 @@ mp_vm_return_kind_t mp_obj_gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_
     mp_obj_gen_instance_t *self = MP_OBJ_TO_PTR(self_in);
     if (self->code_state.ip == 0) {
         // Trying to resume an already stopped generator.
+        #if MICROPY_PY_ASYNC_AWAIT
+        if (self->base.type == &mp_type_agen_instance) {
+            // Do a "raise StopAsyncIteration()".
+            *ret_val = mp_obj_new_exception(&mp_type_StopAsyncIteration);
+            return MP_VM_RETURN_EXCEPTION;
+        }
+        #endif
         // This is an optimised "raise StopIteration(None)".
         *ret_val = mp_const_none;
         return MP_VM_RETURN_NORMAL;
@@ -235,13 +254,42 @@ mp_vm_return_kind_t mp_obj_gen_resume(mp_obj_t self_in, mp_obj_t send_value, mp_
             self->code_state.ip = 0;
             // This is an optimised "raise StopIteration(*ret_val)".
             *ret_val = *self->code_state.sp;
+            #if MICROPY_PY_ASYNC_AWAIT
+            if (self->base.type == &mp_type_agen_instance) {
+                // Reached the end of the async generator, do "raise StopAsyncIteration()".
+                assert(*ret_val == mp_const_none); // enforced by the compiler
+                *ret_val = mp_obj_new_exception(&mp_type_StopAsyncIteration);
+                ret_kind = MP_VM_RETURN_EXCEPTION;
+            }
+            #endif
             break;
 
         case MP_VM_RETURN_YIELD:
+        case MP_VM_RETURN_YIELD_FROM:
             *ret_val = *self->code_state.sp;
             #if MICROPY_PY_GENERATOR_PEND_THROW
             *self->code_state.sp = mp_const_none;
             #endif
+            #if MICROPY_PY_ASYNC_AWAIT
+            if (self->base.type == &mp_type_agen_instance) {
+                #if MICROPY_EMIT_NATIVE
+                if (self->code_state.exc_sp_idx == MP_CODE_STATE_EXC_SP_IDX_SENTINEL) {
+                    // A native generator.
+                    if (ret_kind == MP_VM_RETURN_YIELD) {
+                        // "yield" in native async generator.
+                        ret_kind = MP_VM_RETURN_NORMAL;
+                    } else {
+                        // "yield from" in native async generator.
+                        ret_kind = MP_VM_RETURN_YIELD;
+                    }
+                } else
+                #endif
+                if (*self->code_state.ip != MP_BC_YIELD_FROM) {
+                    // "yield" in async generator.
+                    ret_kind = MP_VM_RETURN_NORMAL;
+                }
+            }
+            #endif
             break;
 
         case MP_VM_RETURN_EXCEPTION: {
@@ -374,3 +422,25 @@ MP_DEFINE_CONST_OBJ_TYPE(
     iter, gen_instance_iternext,
     locals_dict, &gen_instance_locals_dict
     );
+
+/******************************************************************************/
+// async generator instance
+
+#if MICROPY_PY_ASYNC_AWAIT
+
+STATIC const mp_rom_map_elem_t agen_instance_locals_dict_table[] = {
+    { MP_ROM_QSTR(MP_QSTR___aiter__), MP_ROM_PTR(&mp_identity_obj) },
+    { MP_ROM_QSTR(MP_QSTR___anext__), MP_ROM_PTR(&mp_identity_obj) },
+};
+STATIC MP_DEFINE_CONST_DICT(agen_instance_locals_dict, agen_instance_locals_dict_table);
+
+MP_DEFINE_CONST_OBJ_TYPE(
+    mp_type_agen_instance,
+    MP_QSTR_async_generator,
+    MP_TYPE_FLAG_ITER_IS_ITERNEXT,
+    unary_op, mp_generic_unary_op,
+    iter, gen_instance_iternext,
+    locals_dict, &agen_instance_locals_dict
+    );
+
+#endif // MICROPY_PY_ASYNC_AWAIT
diff --git a/py/runtime.c b/py/runtime.c
@@ -1387,7 +1387,11 @@ mp_vm_return_kind_t mp_resume(mp_obj_t self_in, mp_obj_t send_value, mp_obj_t th
     assert((send_value != MP_OBJ_NULL) ^ (throw_value != MP_OBJ_NULL));
     const mp_obj_type_t *type = mp_obj_get_type(self_in);
 
-    if (type == &mp_type_gen_instance) {
+    if (type == &mp_type_gen_instance
+        #if MICROPY_PY_ASYNC_AWAIT
+        || type == &mp_type_agen_instance
+        #endif
+        ) {
         return mp_obj_gen_resume(self_in, send_value, throw_value, ret_val);
     }
 

diff --git a/py/runtime.h b/py/runtime.h
@@ -32,6 +32,7 @@
 typedef enum {
     MP_VM_RETURN_NORMAL,
     MP_VM_RETURN_YIELD,
+    MP_VM_RETURN_YIELD_FROM, // used only by the native emitter
     MP_VM_RETURN_EXCEPTION,
 } mp_vm_return_kind_t;
 

diff --git a/py/runtime0.h b/py/runtime0.h
@@ -26,19 +26,23 @@
 #ifndef MICROPY_INCLUDED_PY_RUNTIME0_H
 #define MICROPY_INCLUDED_PY_RUNTIME0_H
 
-// The first four must fit in 8 bits, see emitbc.c
-// The remaining must fit in 16 bits, see scope.h
-#define MP_SCOPE_FLAG_ALL_SIG      (0x0f)
+// The first five must fit in 7 bits, see mp_raw_code_t.scope_flags.
+// The remaining must fit in 16 bits, see scope_t.scope_flags.
+#define MP_SCOPE_FLAG_ALL_SIG      (0x1f)
+
 #define MP_SCOPE_FLAG_GENERATOR    (0x01)
 #define MP_SCOPE_FLAG_VARKEYWORDS  (0x02)
 #define MP_SCOPE_FLAG_VARARGS      (0x04)
 #define MP_SCOPE_FLAG_DEFKWARGS    (0x08)
-#define MP_SCOPE_FLAG_REFGLOBALS   (0x10) // used only if native emitter enabled
-#define MP_SCOPE_FLAG_HASCONSTS    (0x20) // used only if native emitter enabled
-#define MP_SCOPE_FLAG_VIPERRET_POS    (6) // 3 bits used for viper return type, to pass from compiler to native emitter
-#define MP_SCOPE_FLAG_VIPERRELOC   (0x10) // used only when loading viper from .mpy
-#define MP_SCOPE_FLAG_VIPERRODATA  (0x20) // used only when loading viper from .mpy
-#define MP_SCOPE_FLAG_VIPERBSS     (0x40) // used only when loading viper from .mpy
+#define MP_SCOPE_FLAG_ASYNCGENERATOR (0x10)
+
+#define MP_SCOPE_FLAG_REFGLOBALS   (0x20) // used only if native emitter enabled
+#define MP_SCOPE_FLAG_HASCONSTS    (0x40) // used only if native emitter enabled
+#define MP_SCOPE_FLAG_VIPERRET_POS    (7) // 3 bits used for viper return type, to pass from compiler to native emitter
+#define MP_SCOPE_FLAG_VIPERRELOC   (0x20) // used only when loading viper from .mpy
+#define MP_SCOPE_FLAG_VIPERRODATA  (0x40) // used only when loading viper from .mpy
+#define MP_SCOPE_FLAG_VIPERBSS     (0x80) // used only when loading viper from .mpy
+#define MP_SCOPE_FLAG_ASYNCDEF     (0x8000) // used ony by the compiler
 
 // types for native (viper) function signature
 #define MP_NATIVE_TYPE_OBJ  (0x00)

diff --git a/tests/basics/async_agen.py b/tests/basics/async_agen.py
@@ -0,0 +1,79 @@
+# Test async generators
+
+# helper
+def run_task(t):
+    while True:
+        try:
+            t.send(None)
+        except Exception as er:
+            print('run_task finished with', repr(er))
+            break
+
+# Test SyntaxError
+
+try:
+    exec('async def f(): yield from y')
+except Exception as er:
+    print(type(er))
+
+try:
+    exec('async def f(): yield; return 1')
+except Exception as er:
+    print(type(er))
+
+try:
+    exec('async def f(): return 2; yield')
+except Exception as er:
+    print(type(er))
+
+# Test core behaviour
+
+async def genfunc():
+    yield 1
+    yield 2
+
+async def main():
+    gen = genfunc()
+
+    print(gen.__aiter__() is gen)
+
+    print(repr(await gen.__anext__())) # should be 1
+    print(repr(await gen.__anext__())) # should be 2
+
+    try:
+        await gen.__anext__()
+    except Exception as er:
+        print(type(er))
+
+    try:
+        await gen.__anext__()
+    except Exception as er:
+        print(type(er))
+
+run_task(main())
+
+# Test async-for iterating over an async generator
+
+class Wait:
+    def __init__(self, n):
+        self.n = n
+    def __await__(self):
+        return self
+    __iter__ = __await__ # needed for uPy
+    def __next__(self):
+        if self.n == 0:
+            raise StopIteration
+        print('wait', self.n)
+        self.n -= 1
+
+async def gen():
+    for i in range(4):
+        yield i
+        await Wait(2)
+
+async def main():
+    g = gen()
+    async for a in g:
+        print('got', a)
+
+run_task(main())