gh-121637: Syntax error for optimized-away incorrect await (#121656)

Co-authored-by: Bénédikt Tran <10796600+picnixz@users.noreply.github.com>

Author: JelleZijlstra

Doc/whatsnew/3.14.rst

@@ -75,6 +75,11 @@ New Features
 Other Language Changes
 ======================
 
+* Incorrect usage of :keyword:`await` and asynchronous comprehensions
+  is now detected even if the code is optimized away by the :option:`-O`
+  command line option. For example, ``python -O -c 'assert await 1'``
+  now produces a :exc:`SyntaxError`. (Contributed by Jelle Zijlstra in :gh:`121637`.)
+
 * Added class methods :meth:`float.from_number` and :meth:`complex.from_number`
   to convert a number to :class:`float` or :class:`complex` type correspondingly.
   They raise an error if the argument is a string.

Lib/test/test_builtin.py

@@ -16,6 +16,7 @@
 import random
 import re
 import sys
+import textwrap
 import traceback
 import types
 import typing
@@ -412,7 +413,7 @@ def test_compile_top_level_await_no_coro(self):
         "socket.accept is broken"
     )
     def test_compile_top_level_await(self):
-        """Test whether code some top level await can be compiled.
+        """Test whether code with top level await can be compiled.
 
         Make sure it compiles only with the PyCF_ALLOW_TOP_LEVEL_AWAIT flag
         set, and make sure the generated code object has the CO_COROUTINE flag
@@ -426,6 +427,7 @@ async def arange(n):
                 yield i
 
         modes = ('single', 'exec')
+        optimizations = (-1, 0, 1, 2)
         code_samples = [
             '''a = await asyncio.sleep(0, result=1)''',
             '''async for i in arange(1):
@@ -438,34 +440,52 @@ async def arange(n):
             '''a = [x async for x in arange(2) async for x in arange(2)][1]''',
             '''a = [x async for x in (x async for x in arange(5))][1]''',
             '''a, = [1 for x in {x async for x in arange(1)}]''',
-            '''a = [await asyncio.sleep(0, x) async for x in arange(2)][1]'''
+            '''a = [await asyncio.sleep(0, x) async for x in arange(2)][1]''',
+            # gh-121637: Make sure we correctly handle the case where the
+            # async code is optimized away
+            '''assert not await asyncio.sleep(0); a = 1''',
+            '''assert [x async for x in arange(1)]; a = 1''',
+            '''assert {x async for x in arange(1)}; a = 1''',
+            '''assert {x: x async for x in arange(1)}; a = 1''',
+            '''
+            if (a := 1) and __debug__:
+                async with asyncio.Lock() as l:
+                    pass
+            ''',
+            '''
+            if (a := 1) and __debug__:
+                async for x in arange(2):
+                    pass
+            ''',
         ]
         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)
-
-                co = compile(source,
-                             '?',
-                             mode,
-                             flags=ast.PyCF_ALLOW_TOP_LEVEL_AWAIT)
-
-                self.assertEqual(co.co_flags & CO_COROUTINE, CO_COROUTINE,
-                                 msg=f"source={source} mode={mode}")
+            for mode, code_sample, optimize in product(modes, code_samples, optimizations):
+                with self.subTest(mode=mode, code_sample=code_sample, optimize=optimize):
+                    source = dedent(code_sample)
+                    with self.assertRaises(
+                            SyntaxError, msg=f"source={source} mode={mode}"):
+                        compile(source, '?', mode, optimize=optimize)
 
-                # test we can create and  advance a function type
-                globals_ = {'asyncio': asyncio, 'a': 0, 'arange': arange}
-                async_f = FunctionType(co, globals_)
-                asyncio.run(async_f())
-                self.assertEqual(globals_['a'], 1)
-
-                # test we can await-eval,
-                globals_ = {'asyncio': asyncio, 'a': 0, 'arange': arange}
-                asyncio.run(eval(co, globals_))
-                self.assertEqual(globals_['a'], 1)
+                    co = compile(source,
+                                '?',
+                                mode,
+                                flags=ast.PyCF_ALLOW_TOP_LEVEL_AWAIT,
+                                optimize=optimize)
+
+                    self.assertEqual(co.co_flags & CO_COROUTINE, CO_COROUTINE,
+                                    msg=f"source={source} mode={mode}")
+
+                    # test we can create and  advance a function type
+                    globals_ = {'asyncio': asyncio, 'a': 0, 'arange': arange}
+                    async_f = FunctionType(co, globals_)
+                    asyncio.run(async_f())
+                    self.assertEqual(globals_['a'], 1)
+
+                    # test we can await-eval,
+                    globals_ = {'asyncio': asyncio, 'a': 0, 'arange': arange}
+                    asyncio.run(eval(co, globals_))
+                    self.assertEqual(globals_['a'], 1)
         finally:
             asyncio.set_event_loop_policy(policy)
 

Misc/NEWS.d/next/Core and Builtins/2024-07-12-18-18-44.gh-issue-121297.67VE7b.rst

@@ -0,0 +1,4 @@
+Previously, incorrect usage of :keyword:`await` or asynchronous
+comprehensions in code removed by the :option:`-O` option was not flagged by
+the Python compiler. Now, such code raises :exc:`SyntaxError`. Patch by
+Jelle Zijlstra.

Python/compile.c

@@ -5675,14 +5675,16 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
     PyCodeObject *co = NULL;
     inlined_comprehension_state inline_state = {NULL, NULL, NULL, NO_LABEL, NO_LABEL};
     comprehension_ty outermost;
+#ifndef NDEBUG
     int scope_type = c->u->u_scope_type;
     int is_top_level_await = IS_TOP_LEVEL_AWAIT(c);
+#endif
     PySTEntryObject *entry = _PySymtable_Lookup(SYMTABLE(c), (void *)e);
     if (entry == NULL) {
         goto error;
     }
     int is_inlined = entry->ste_comp_inlined;
-    int is_async_generator = entry->ste_coroutine;
+    int is_async_comprehension = entry->ste_coroutine;
 
     location loc = LOC(e);
 
@@ -5697,22 +5699,17 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
     }
     else {
         if (compiler_enter_scope(c, name, COMPILER_SCOPE_COMPREHENSION,
-                                (void *)e, e->lineno, NULL) < 0)
-        {
+                                (void *)e, e->lineno, NULL) < 0) {
             goto error;
         }
     }
     Py_CLEAR(entry);
 
-    if (is_async_generator && type != COMP_GENEXP &&
-        scope_type != COMPILER_SCOPE_ASYNC_FUNCTION &&
-        scope_type != COMPILER_SCOPE_COMPREHENSION &&
-        !is_top_level_await)
-    {
-        compiler_error(c, loc, "asynchronous comprehension outside of "
-                               "an asynchronous function");
-        goto error_in_scope;
-    }
+    assert (!is_async_comprehension ||
+            type == COMP_GENEXP ||
+            scope_type == COMPILER_SCOPE_ASYNC_FUNCTION ||
+            scope_type == COMPILER_SCOPE_COMPREHENSION ||
+            is_top_level_await);
 
     if (type != COMP_GENEXP) {
         int op;
@@ -5777,7 +5774,7 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
 
     ADDOP_I(c, loc, CALL, 0);
 
-    if (is_async_generator && type != COMP_GENEXP) {
+    if (is_async_comprehension && type != COMP_GENEXP) {
         ADDOP_I(c, loc, GET_AWAITABLE, 0);
         ADDOP_LOAD_CONST(c, loc, Py_None);
         ADD_YIELD_FROM(c, loc, 1);
@@ -6138,16 +6135,10 @@ compiler_visit_expr(struct compiler *c, expr_ty e)
         ADD_YIELD_FROM(c, loc, 0);
         break;
     case Await_kind:
-        if (!IS_TOP_LEVEL_AWAIT(c)){
-            if (!_PyST_IsFunctionLike(SYMTABLE_ENTRY(c))) {
-                return compiler_error(c, loc, "'await' outside function");
-            }
-
-            if (c->u->u_scope_type != COMPILER_SCOPE_ASYNC_FUNCTION &&
-                    c->u->u_scope_type != COMPILER_SCOPE_COMPREHENSION) {
-                return compiler_error(c, loc, "'await' outside async function");
-            }
-        }
+        assert(IS_TOP_LEVEL_AWAIT(c) || (_PyST_IsFunctionLike(SYMTABLE_ENTRY(c)) && (
+            c->u->u_scope_type == COMPILER_SCOPE_ASYNC_FUNCTION ||
+            c->u->u_scope_type == COMPILER_SCOPE_COMPREHENSION
+        )));
 
         VISIT(c, expr, e->v.Await.value);
         ADDOP_I(c, loc, GET_AWAITABLE, 0);

Python/symtable.c

@@ -70,10 +70,10 @@
 #define DUPLICATE_TYPE_PARAM \
 "duplicate type parameter '%U'"
 
-#define ASYNC_WITH_OUTISDE_ASYNC_FUNC \
+#define ASYNC_WITH_OUTSIDE_ASYNC_FUNC \
 "'async with' outside async function"
 
-#define ASYNC_FOR_OUTISDE_ASYNC_FUNC \
+#define ASYNC_FOR_OUTSIDE_ASYNC_FUNC \
 "'async for' outside async function"
 
 #define LOCATION(x) SRC_LOCATION_FROM_AST(x)
@@ -82,6 +82,8 @@
     PyErr_RangedSyntaxLocationObject((FNAME), \
         (L).lineno, (L).col_offset + 1, (L).end_lineno, (L).end_col_offset + 1)
 
+#define IS_ASYNC_DEF(st) ((st)->st_cur->ste_type == FunctionBlock && (st)->st_cur->ste_coroutine)
+
 static PySTEntryObject *
 ste_new(struct symtable *st, identifier name, _Py_block_ty block,
         void *key, _Py_SourceLocation loc)
@@ -1660,12 +1662,18 @@ check_import_from(struct symtable *st, stmt_ty s)
     return 1;
 }
 
+static bool
+allows_top_level_await(struct symtable *st)
+{
+    return (st->st_future->ff_features & PyCF_ALLOW_TOP_LEVEL_AWAIT) &&
+            st->st_cur->ste_type == ModuleBlock;
+}
+
+
 static void
 maybe_set_ste_coroutine_for_module(struct symtable *st, stmt_ty s)
 {
-    if ((st->st_future->ff_features & PyCF_ALLOW_TOP_LEVEL_AWAIT) &&
-        (st->st_cur->ste_type == ModuleBlock))
-    {
+    if (allows_top_level_await(st)) {
         st->st_cur->ste_coroutine = 1;
     }
 }
@@ -2054,15 +2062,15 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
     }
     case AsyncWith_kind:
         maybe_set_ste_coroutine_for_module(st, s);
-        if (!symtable_raise_if_not_coroutine(st, ASYNC_WITH_OUTISDE_ASYNC_FUNC, LOCATION(s))) {
+        if (!symtable_raise_if_not_coroutine(st, ASYNC_WITH_OUTSIDE_ASYNC_FUNC, LOCATION(s))) {
             VISIT_QUIT(st, 0);
         }
         VISIT_SEQ(st, withitem, s->v.AsyncWith.items);
         VISIT_SEQ(st, stmt, s->v.AsyncWith.body);
         break;
     case AsyncFor_kind:
         maybe_set_ste_coroutine_for_module(st, s);
-        if (!symtable_raise_if_not_coroutine(st, ASYNC_FOR_OUTISDE_ASYNC_FUNC, LOCATION(s))) {
+        if (!symtable_raise_if_not_coroutine(st, ASYNC_FOR_OUTSIDE_ASYNC_FUNC, LOCATION(s))) {
             VISIT_QUIT(st, 0);
         }
         VISIT(st, expr, s->v.AsyncFor.target);
@@ -2279,6 +2287,20 @@ symtable_visit_expr(struct symtable *st, expr_ty e)
         if (!symtable_raise_if_annotation_block(st, "await expression", e)) {
             VISIT_QUIT(st, 0);
         }
+        if (!allows_top_level_await(st)) {
+            if (!_PyST_IsFunctionLike(st->st_cur)) {
+                PyErr_SetString(PyExc_SyntaxError,
+                                "'await' outside function");
+                SET_ERROR_LOCATION(st->st_filename, LOCATION(e));
+                VISIT_QUIT(st, 0);
+            }
+            if (!IS_ASYNC_DEF(st) && st->st_cur->ste_comprehension == NoComprehension) {
+                PyErr_SetString(PyExc_SyntaxError,
+                                "'await' outside async function");
+                SET_ERROR_LOCATION(st->st_filename, LOCATION(e));
+                VISIT_QUIT(st, 0);
+            }
+        }
         VISIT(st, expr, e->v.Await.value);
         st->st_cur->ste_coroutine = 1;
         break;
@@ -2798,6 +2820,16 @@ symtable_handle_comprehension(struct symtable *st, expr_ty e,
     if (!symtable_exit_block(st)) {
         return 0;
     }
+    if (is_async &&
+        !IS_ASYNC_DEF(st) &&
+        st->st_cur->ste_comprehension == NoComprehension &&
+        !allows_top_level_await(st))
+    {
+        PyErr_SetString(PyExc_SyntaxError, "asynchronous comprehension outside of "
+                                           "an asynchronous function");
+        SET_ERROR_LOCATION(st->st_filename, LOCATION(e));
+        return 0;
+    }
     if (is_async) {
         st->st_cur->ste_coroutine = 1;
     }