bpo-39965: Correctly raise SyntaxError if await is used outside async functions when PyCF_ALLOW_TOP_LEVEL_AWAIT is set (GH-19010)

3 files changed, 49 insertions, 4 deletions

diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py
index e50c273..d6b9ee1 100644
--- a/Lib/test/test_builtin.py
+++ b/Lib/test/test_builtin.py
@@ -421,6 +421,44 @@ class BuiltinTest(unittest.TestCase):
         finally:
             asyncio.set_event_loop_policy(policy)
 
+    def test_compile_top_level_await_invalid_cases(self):
+         # helper function just to check we can run top=level async-for
+        async def arange(n):
+            for i in range(n):
+                yield i
+
+        modes = ('single', 'exec')
+        code_samples = [
+            '''def f():  await arange(10)\n''',
+            '''def f():  [x async for x in arange(10)]\n''',
+            '''def f():  [await x async for x in arange(10)]\n''',
+            '''def f():
+                   async for i in arange(1):
+                       a = 1
+            ''',
+            '''def f():
+                   async with asyncio.Lock() as l:
+                       a = 1
+            '''
+        ]
+        policy = maybe_get_event_loop_policy()
+        try:
+            for mode, code_sample in product(modes, code_samples):
+                source = dedent(code_sample)
+                with self.assertRaises(
+                        SyntaxError, msg=f"source={source} mode={mode}"):
+                    compile(source, '?', mode)
+
+                with self.assertRaises(
+                        SyntaxError, msg=f"source={source} mode={mode}"):
+                    co = compile(source,
+                             '?',
+                             mode,
+                             flags=ast.PyCF_ALLOW_TOP_LEVEL_AWAIT)
+        finally:
+            asyncio.set_event_loop_policy(policy)
+
+
     def test_compile_async_generator(self):
         """
         With the PyCF_ALLOW_TOP_LEVEL_AWAIT flag added in 3.8, we want to
diff --git a/Misc/NEWS.d/next/Core and Builtins/2020-03-15-03-52-01.bpo-39965.Od3ZdP.rst b/Misc/NEWS.d/next/Core and Builtins/2020-03-15-03-52-01.bpo-39965.Od3ZdP.rst
new file mode 100644
index 0000000..4e3ac7c
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2020-03-15-03-52-01.bpo-39965.Od3ZdP.rst
@@ -0,0 +1,3 @@
+Correctly raise ``SyntaxError`` if *await* is used inside non-async
+functions and ``PyCF_ALLOW_TOP_LEVEL_AWAIT`` is set (like in the asyncio
+REPL). Patch by Pablo Galindo.
diff --git a/Python/compile.c b/Python/compile.c
index 3b4cdaa..b92cb44 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -41,6 +41,10 @@
 #define COMP_SETCOMP  2
 #define COMP_DICTCOMP 3
 
+#define IS_TOP_LEVEL_AWAIT(c) ( \
+        (c->c_flags->cf_flags & PyCF_ALLOW_TOP_LEVEL_AWAIT) \
+        && (c->u->u_ste->ste_type == ModuleBlock))
+
 struct instr {
     unsigned i_jabs : 1;
     unsigned i_jrel : 1;
@@ -2743,7 +2747,7 @@ static int
 compiler_async_for(struct compiler *c, stmt_ty s)
 {
     basicblock *start, *except, *end;
-    if (c->c_flags->cf_flags & PyCF_ALLOW_TOP_LEVEL_AWAIT){
+    if (IS_TOP_LEVEL_AWAIT(c)){
         c->u->u_ste->ste_coroutine = 1;
     } else if (c->u->u_scope_type != COMPILER_SCOPE_ASYNC_FUNCTION) {
         return compiler_error(c, "'async for' outside async function");
@@ -4789,7 +4793,7 @@ compiler_async_with(struct compiler *c, stmt_ty s, int pos)
     withitem_ty item = asdl_seq_GET(s->v.AsyncWith.items, pos);
 
     assert(s->kind == AsyncWith_kind);
-    if (c->c_flags->cf_flags & PyCF_ALLOW_TOP_LEVEL_AWAIT){
+    if (IS_TOP_LEVEL_AWAIT(c)){
         c->u->u_ste->ste_coroutine = 1;
     } else if (c->u->u_scope_type != COMPILER_SCOPE_ASYNC_FUNCTION){
         return compiler_error(c, "'async with' outside async function");
@@ -5007,7 +5011,7 @@ compiler_visit_expr1(struct compiler *c, expr_ty e)
         ADDOP(c, YIELD_FROM);
         break;
     case Await_kind:
-        if (!(c->c_flags->cf_flags & PyCF_ALLOW_TOP_LEVEL_AWAIT)){
+        if (!IS_TOP_LEVEL_AWAIT(c)){
             if (c->u->u_ste->ste_type != FunctionBlock){
                 return compiler_error(c, "'await' outside function");
             }
@@ -5836,7 +5840,7 @@ compute_code_flags(struct compiler *c)
     /* (Only) inherit compilerflags in PyCF_MASK */
     flags |= (c->c_flags->cf_flags & PyCF_MASK);
 
-    if ((c->c_flags->cf_flags & PyCF_ALLOW_TOP_LEVEL_AWAIT) &&
+    if ((IS_TOP_LEVEL_AWAIT(c)) &&
          ste->ste_coroutine &&
          !ste->ste_generator) {
         flags |= CO_COROUTINE;