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

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
 ------

Lib/test/test_compile.py

@@ -876,6 +876,10 @@ class TestSpecifics(unittest.TestCase):
             def with_docstring():
                 "docstring"
 
+            def two_strings():
+                "docstring"
+                "not docstring"
+
             def with_fstring():
                 f"not docstring"
 
@@ -891,8 +895,10 @@ class TestSpecifics(unittest.TestCase):
 
                 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__)
 

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.

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);

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);