Skip to content

gh-123958: move assert optimization from codegen to ast_opt #124143

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 4 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions Doc/whatsnew/3.14.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
------
Expand Down
3 changes: 2 additions & 1 deletion Include/internal/pycore_compile.h
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
29 changes: 23 additions & 6 deletions Lib/test/test_builtin.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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()))
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks like we no longer raise an error in optimized mode if there is an invalid await. That seems wrong and doesn't align with what we decided on in #121637.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we want to stick with this then we probably can't optimise the assertions in ast_opt.

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):
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Assertions are now removed from an optimized AST in optimization level > 0.y
33 changes: 32 additions & 1 deletion Python/ast_opt.c
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
typedef struct {
int optimize;
int ff_features;
PyObject *filename;

int recursion_depth; /* current recursion depth */
int recursion_limit; /* recursion limit */
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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_);
Comment on lines +1035 to +1036
Copy link
Member

@Eclips4 Eclips4 Sep 17, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Technically speaking, this is not the removal of the assert statement; it's replacement of assert with pass statement. Is it possible to completely remove the assert node? Or will it break some existing code?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pass becomes a NOP, which gets removed later. This is how we get rid of code. However, I did forget to wipe out the location information (a NOP is not removed if it's the only instruction from its line in the source code).

}
break;
case Expr_kind:
CALL(astfold_expr, expr_ty, node_->v.Expr.value);
Expand Down Expand Up @@ -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();
Expand Down
15 changes: 1 addition & 14 deletions Python/codegen.c
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Left this here to convince ourselves of correctness. I'll remove it before committing.

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);
Expand Down
4 changes: 2 additions & 2 deletions Python/compile.c
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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;
Expand Down
Loading