From 2b8e88c563e4bab53e3183d15f04795d645eea28 Mon Sep 17 00:00:00 2001
From: Damien George <damien@micropython.org>
Date: Sat, 24 Jun 2023 17:02:58 +1000
Subject: [PATCH] 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 <damien@micropython.org>
---
 py/compile.c  | 13 +++++++++++--
 py/compile.h  |  5 +++++
 py/mpconfig.h |  5 +++++
 3 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/py/compile.c b/py/compile.c
index 7a359e662e..62757de3c0 100644
--- 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
index 5e0fd8b31c..f9970a521d 100644
--- 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
index d9cff930d1..af2480266b 100644
--- 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)