py/compile: Add option to allow compiling top-level await.

Enabled by MICROPY_COMPILE_ALLOW_TOP_LEVEL_AWAIT. When enabled, this means that scope such as module-level functions and REPL statements can yield. The outer C code must then handle this yielded generator. Signed-off-by: Damien George <[email protected]>
picoscratch · Mar 22, 2024 · 2b8e88c · 2b8e88c
1 parent acbdbcd
commit 2b8e88c
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 2 deletions.
diff --git a/py/compile.c b/py/compile.c
@@ -196,6 +196,10 @@ typedef struct _compiler_t {
     mp_emit_common_t emit_common;
 } compiler_t;
 
+#if MICROPY_COMP_ALLOW_TOP_LEVEL_AWAIT
+bool mp_compile_allow_top_level_await = false;
+#endif
+
 /******************************************************************************/
 // mp_emit_common_t helper functions
 // These are defined here so they can be inlined, to reduce code size.
@@ -2759,8 +2763,13 @@ static void compile_yield_expr(compiler_t *comp, mp_parse_node_struct_t *pns) {
 #if MICROPY_PY_ASYNC_AWAIT
 static void compile_atom_expr_await(compiler_t *comp, mp_parse_node_struct_t *pns) {
     if (comp->scope_cur->kind != SCOPE_FUNCTION && comp->scope_cur->kind != SCOPE_LAMBDA) {
-        compile_syntax_error(comp, (mp_parse_node_t)pns, MP_ERROR_TEXT("'await' outside function"));
-        return;
+        #if MICROPY_COMP_ALLOW_TOP_LEVEL_AWAIT
+        if (!mp_compile_allow_top_level_await)
+        #endif
+        {
+            compile_syntax_error(comp, (mp_parse_node_t)pns, MP_ERROR_TEXT("'await' outside function"));
+            return;
+        }
     }
     compile_atom_expr_normal(comp, pns);
     compile_yield_from(comp);

diff --git a/py/compile.h b/py/compile.h
@@ -30,6 +30,11 @@
 #include "py/parse.h"
 #include "py/emitglue.h"
 
+#if MICROPY_COMP_ALLOW_TOP_LEVEL_AWAIT
+// set to `true` to allow top-level await expressions
+extern bool mp_compile_allow_top_level_await;
+#endif
+
 // the compiler will raise an exception if an error occurred
 // the compiler will clear the parse tree before it returns
 // mp_globals_get() will be used for the context

diff --git a/py/mpconfig.h b/py/mpconfig.h
@@ -446,6 +446,11 @@
 #define MICROPY_DYNAMIC_COMPILER (0)
 #endif
 
+// Whether the compiler allows compiling top-level await expressions
+#ifndef MICROPY_COMP_ALLOW_TOP_LEVEL_AWAIT
+#define MICROPY_COMP_ALLOW_TOP_LEVEL_AWAIT (0)
+#endif
+
 // Whether to enable constant folding; eg 1+2 rewritten as 3
 #ifndef MICROPY_COMP_CONST_FOLDING
 #define MICROPY_COMP_CONST_FOLDING (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_CORE_FEATURES)