From e07154fd1e3152a758cf9b476257a4ffdc48dfc6 Mon Sep 17 00:00:00 2001 From: Irit Katriel <1055913+iritkatriel@users.noreply.github.com> Date: Wed, 11 Sep 2024 18:02:28 +0100 Subject: [PATCH] gh-123958: apply docstring removal optimization in ast_opt instead of codegen (#123959) --- Doc/whatsnew/3.14.rst | 4 +- Lib/test/test_compile.py | 6 +++ ...-09-11-15-48-36.gh-issue-123958.5VW2r0.rst | 1 + Python/ast_opt.c | 21 ++++++++++ Python/codegen.c | 38 ++++++++----------- 5 files changed, 47 insertions(+), 23 deletions(-) create mode 100644 Misc/NEWS.d/next/Core_and_Builtins/2024-09-11-15-48-36.gh-issue-123958.5VW2r0.rst diff --git a/Doc/whatsnew/3.14.rst b/Doc/whatsnew/3.14.rst index 7aca6bd2117173..dc2083b241d391 100644 --- a/Doc/whatsnew/3.14.rst +++ b/Doc/whatsnew/3.14.rst @@ -123,9 +123,11 @@ ast (Contributed by Batuhan Taskaya and Jeremy Hylton in :issue:`15987`.) * Add support for :func:`copy.replace` for AST nodes. - (Contributed by Bénédikt Tran in :gh:`121141`.) +* Docstrings are now removed from an optimized AST in optimization level 2. + (Contributed by Irit Katriel in :gh:`123958`.) + ctypes ------ diff --git a/Lib/test/test_compile.py b/Lib/test/test_compile.py index 7d6ddba1adc87e..f22761f0a3af9f 100644 --- a/Lib/test/test_compile.py +++ b/Lib/test/test_compile.py @@ -876,6 +876,10 @@ def test_docstring(self): def with_docstring(): "docstring" + def two_strings(): + "docstring" + "not docstring" + def with_fstring(): f"not docstring" @@ -891,8 +895,10 @@ def with_const_expression(): if opt < 2: self.assertEqual(ns['with_docstring'].__doc__, "docstring") + self.assertEqual(ns['two_strings'].__doc__, "docstring") else: self.assertIsNone(ns['with_docstring'].__doc__) + self.assertIsNone(ns['two_strings'].__doc__) self.assertIsNone(ns['with_fstring'].__doc__) self.assertIsNone(ns['with_const_expression'].__doc__) diff --git a/Misc/NEWS.d/next/Core_and_Builtins/2024-09-11-15-48-36.gh-issue-123958.5VW2r0.rst b/Misc/NEWS.d/next/Core_and_Builtins/2024-09-11-15-48-36.gh-issue-123958.5VW2r0.rst new file mode 100644 index 00000000000000..fc2623a7cbf789 --- /dev/null +++ b/Misc/NEWS.d/next/Core_and_Builtins/2024-09-11-15-48-36.gh-issue-123958.5VW2r0.rst @@ -0,0 +1 @@ +docstrings are now removed from the optimized AST in optimization level 2. diff --git a/Python/ast_opt.c b/Python/ast_opt.c index 5a51305d2a7ade..f5b04757e08bf3 100644 --- a/Python/ast_opt.c +++ b/Python/ast_opt.c @@ -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); diff --git a/Python/codegen.c b/Python/codegen.c index ed06724e95d05b..fd2260a538f76d 100644 --- a/Python/codegen.c +++ b/Python/codegen.c @@ -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);