gh-123958: apply docstring removal optimization in ast_opt instead of codegen (#123959)

This commit is contained in:
Irit Katriel 2024-09-11 18:02:28 +01:00 committed by GitHub
parent 2938c3dec9
commit e07154fd1e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 47 additions and 23 deletions

View file

@ -673,10 +673,31 @@ static int astfold_type_param(type_param_ty node_, PyArena *ctx_, _PyASTOptimize
}
static int
stmt_seq_remove_item(asdl_stmt_seq *stmts, Py_ssize_t idx)
{
if (idx >= asdl_seq_LEN(stmts)) {
return 0;
}
for (Py_ssize_t i = idx; i < asdl_seq_LEN(stmts) - 1; i++) {
stmt_ty st = (stmt_ty)asdl_seq_GET(stmts, i+1);
asdl_seq_SET(stmts, i, st);
}
stmts->size--;
return 1;
}
static int
astfold_body(asdl_stmt_seq *stmts, PyArena *ctx_, _PyASTOptimizeState *state)
{
int docstring = _PyAST_GetDocString(stmts) != NULL;
if (docstring && (state->optimize >= 2)) {
/* remove the docstring */
if (!stmt_seq_remove_item(stmts, 0)) {
return 0;
}
docstring = 0;
}
CALL_SEQ(astfold_stmt, stmt, stmts);
if (!docstring && _PyAST_GetDocString(stmts) != NULL) {
stmt_ty st = (stmt_ty)asdl_seq_GET(stmts, 0);

View file

@ -763,19 +763,18 @@ _PyCodegen_Body(compiler *c, location loc, asdl_stmt_seq *stmts)
PyObject *docstring = _PyAST_GetDocString(stmts);
if (docstring) {
first_instr = 1;
/* if not -OO mode, set docstring */
if (OPTIMIZATION_LEVEL(c) < 2) {
PyObject *cleandoc = _PyCompile_CleanDoc(docstring);
if (cleandoc == NULL) {
return ERROR;
}
stmt_ty st = (stmt_ty)asdl_seq_GET(stmts, 0);
assert(st->kind == Expr_kind);
location loc = LOC(st->v.Expr.value);
ADDOP_LOAD_CONST(c, loc, cleandoc);
Py_DECREF(cleandoc);
RETURN_IF_ERROR(codegen_nameop(c, NO_LOCATION, &_Py_ID(__doc__), Store));
/* set docstring */
assert(OPTIMIZATION_LEVEL(c) < 2);
PyObject *cleandoc = _PyCompile_CleanDoc(docstring);
if (cleandoc == NULL) {
return ERROR;
}
stmt_ty st = (stmt_ty)asdl_seq_GET(stmts, 0);
assert(st->kind == Expr_kind);
location loc = LOC(st->v.Expr.value);
ADDOP_LOAD_CONST(c, loc, cleandoc);
Py_DECREF(cleandoc);
RETURN_IF_ERROR(codegen_nameop(c, NO_LOCATION, &_Py_ID(__doc__), Store));
}
}
for (Py_ssize_t i = first_instr; i < asdl_seq_LEN(stmts); i++) {
@ -1230,18 +1229,13 @@ codegen_function_body(compiler *c, stmt_ty s, int is_async, Py_ssize_t funcflags
Py_ssize_t first_instr = 0;
PyObject *docstring = _PyAST_GetDocString(body);
assert(OPTIMIZATION_LEVEL(c) < 2 || docstring == NULL);
if (docstring) {
first_instr = 1;
/* if not -OO mode, add docstring */
if (OPTIMIZATION_LEVEL(c) < 2) {
docstring = _PyCompile_CleanDoc(docstring);
if (docstring == NULL) {
_PyCompile_ExitScope(c);
return ERROR;
}
}
else {
docstring = NULL;
docstring = _PyCompile_CleanDoc(docstring);
if (docstring == NULL) {
_PyCompile_ExitScope(c);
return ERROR;
}
}
Py_ssize_t idx = _PyCompile_AddConst(c, docstring ? docstring : Py_None);