gh-123958: move assert optimization from codegen to ast_opt

Doc/whatsnew/3.14.rst

@@ -128,6 +128,8 @@ ast
 * Docstrings are now removed from an optimized AST in optimization level 2.
   (Contributed by Irit Katriel in :gh:`123958`.)
 
+* Assertions are now removed from an optimized AST in optimization level > 0.
+  (Contributed by Irit Katriel in :gh:`123958`.)
 
 ctypes
 ------

Include/internal/pycore_compile.h

@@ -37,7 +37,8 @@ extern int _PyAST_Optimize(
     struct _mod *,
     struct _arena *arena,
     int optimize,
-    int ff_features);
+    int ff_features,
+    PyObject *filename);
 
 
 typedef struct {

Lib/test/test_builtin.py

@@ -443,12 +443,6 @@ async def arange(n):
             '''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]''',
-            # 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:
@@ -491,6 +485,29 @@ async def arange(n):
         finally:
             asyncio.set_event_loop_policy(policy)
 
+    def test_top_level_async_in_assert_compiles(self):
+        # gh-121637: do the right thing when async code is optimized away
+        modes = ('single', 'exec')
+        optimizations = (0, 1, 2)
+        code_samples = [
+            '''assert not await asyncio.sleep(0)''',
+            '''assert [x async for x in arange(1)]''',
+            '''assert {x async for x in arange(1)}''',
+            '''assert {x: x async for x in arange(1)}''',
+        ]
+        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)
+                if optimize:
+                    co = compile(source, '?', mode, optimize=optimize)
+                    co_expected = compile('pass', '?', mode, optimize=optimize)
+                    self.assertEqual(list(co.co_lines()), list(co_expected.co_lines()))
+                else:
+                    with self.assertRaises(SyntaxError,
+                                           msg=f"source={source} mode={mode}"):
+                        compile(source, '?', mode, optimize=optimize)
+
+
     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):

Misc/NEWS.d/next/Core and Builtins/2024-09-16-22-44-29.gh-issue-123958.THsRVv.rst

@@ -0,0 +1 @@
+Assertions are now removed from an optimized AST in optimization level > 0.y

Python/ast_opt.c

@@ -10,6 +10,7 @@
 typedef struct {
     int optimize;
     int ff_features;
+    PyObject *filename;
 
     int recursion_depth;            /* current recursion depth */
     int recursion_limit;            /* recursion limit */
@@ -51,6 +52,12 @@ make_const(expr_ty node, PyObject *val, PyArena *arena)
     return 1;
 }
 
+static void
+make_pass(stmt_ty node)
+{
+    node->kind = Pass_kind;
+}
+
 #define COPY_NODE(TO, FROM) (memcpy((TO), (FROM), sizeof(struct _expr)))
 
 static int
@@ -1005,6 +1012,29 @@ astfold_stmt(stmt_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
     case Assert_kind:
         CALL(astfold_expr, expr_ty, node_->v.Assert.test);
         CALL_OPT(astfold_expr, expr_ty, node_->v.Assert.msg);
+        /* Always emit a warning if the test is a non-zero length tuple */
+        if ((node_->v.Assert.test->kind == Tuple_kind &&
+            asdl_seq_LEN(node_->v.Assert.test->v.Tuple.elts) > 0) ||
+            (node_->v.Assert.test->kind == Constant_kind &&
+             PyTuple_Check(node_->v.Assert.test->v.Constant.value) &&
+             PyTuple_Size(node_->v.Assert.test->v.Constant.value) > 0))
+        {
+            PyObject *msg = PyUnicode_FromString("assertion is always true, "
+                                                 "perhaps remove parentheses?");
+            if (msg == NULL) {
+                return 0;
+            }
+            int ret = _PyErr_EmitSyntaxWarning(msg, state->filename,
+                                               node_->lineno, node_->col_offset + 1,
+                                               node_->end_lineno, node_->end_col_offset + 1);
+            Py_DECREF(msg);
+            if (ret < 0) {
+                return 0;
+            }
+        }
+        if (state->optimize) {
+            make_pass(node_);
+        }
         break;
     case Expr_kind:
         CALL(astfold_expr, expr_ty, node_->v.Expr.value);
@@ -1125,14 +1155,15 @@ astfold_type_param(type_param_ty node_, PyArena *ctx_, _PyASTOptimizeState *stat
 #undef CALL_SEQ
 
 int
-_PyAST_Optimize(mod_ty mod, PyArena *arena, int optimize, int ff_features)
+_PyAST_Optimize(mod_ty mod, PyArena *arena, int optimize, int ff_features, PyObject *filename)
 {
     PyThreadState *tstate;
     int starting_recursion_depth;
 
     _PyASTOptimizeState state;
     state.optimize = optimize;
     state.ff_features = ff_features;
+    state.filename = filename;
 
     /* Setup recursion depth check counters */
     tstate = _PyThreadState_GET();

Python/codegen.c

@@ -2792,20 +2792,7 @@ codegen_from_import(compiler *c, stmt_ty s)
 static int
 codegen_assert(compiler *c, stmt_ty s)
 {
-    /* Always emit a warning if the test is a non-zero length tuple */
-    if ((s->v.Assert.test->kind == Tuple_kind &&
-        asdl_seq_LEN(s->v.Assert.test->v.Tuple.elts) > 0) ||
-        (s->v.Assert.test->kind == Constant_kind &&
-         PyTuple_Check(s->v.Assert.test->v.Constant.value) &&
-         PyTuple_Size(s->v.Assert.test->v.Constant.value) > 0))
-    {
-        RETURN_IF_ERROR(
-            _PyCompile_Warn(c, LOC(s), "assertion is always true, "
-                                       "perhaps remove parentheses?"));
-    }
-    if (OPTIMIZATION_LEVEL(c)) {
-        return SUCCESS;
-    }
+    assert(!OPTIMIZATION_LEVEL(c));
     NEW_JUMP_TARGET_LABEL(c, end);
     RETURN_IF_ERROR(codegen_jump_if(c, LOC(s), s->v.Assert.test, end, 1));
     ADDOP_I(c, LOC(s), LOAD_COMMON_CONSTANT, CONSTANT_ASSERTIONERROR);

Python/compile.c

@@ -124,7 +124,7 @@ compiler_setup(compiler *c, mod_ty mod, PyObject *filename,
     c->c_optimize = (optimize == -1) ? _Py_GetConfig()->optimization_level : optimize;
     c->c_save_nested_seqs = false;
 
-    if (!_PyAST_Optimize(mod, arena, c->c_optimize, merged)) {
+    if (!_PyAST_Optimize(mod, arena, c->c_optimize, merged, filename)) {
         return ERROR;
     }
     c->c_st = _PySymtable_Build(mod, filename, &c->c_future);
@@ -1381,7 +1381,7 @@ _PyCompile_AstOptimize(mod_ty mod, PyObject *filename, PyCompilerFlags *cf,
     if (optimize == -1) {
         optimize = _Py_GetConfig()->optimization_level;
     }
-    if (!_PyAST_Optimize(mod, arena, optimize, flags)) {
+    if (!_PyAST_Optimize(mod, arena, optimize, flags, filename)) {
         return -1;
     }
     return 0;