gh-124022: Fix bug where class docstring is removed in interactive mode (#124023)

Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>

5 files changed, 37 insertions, 12 deletions

diff --git a/Include/internal/pycore_compile.h b/Include/internal/pycore_compile.h
index 51db6fc..5359f2d 100644
--- a/Include/internal/pycore_compile.h
+++ b/Include/internal/pycore_compile.h
@@ -172,7 +172,8 @@ int _PyCompile_AddDeferredAnnotaion(struct _PyCompiler *c, stmt_ty s);
 int _PyCodegen_AddReturnAtEnd(struct _PyCompiler *c, int addNone);
 int _PyCodegen_EnterAnonymousScope(struct _PyCompiler* c, mod_ty mod);
 int _PyCodegen_Expression(struct _PyCompiler *c, expr_ty e);
-int _PyCodegen_Body(struct _PyCompiler *c, _Py_SourceLocation loc, asdl_stmt_seq *stmts);
+int _PyCodegen_Body(struct _PyCompiler *c, _Py_SourceLocation loc, asdl_stmt_seq *stmts,
+                    bool is_interactive);
 
 /* Utility for a number of growing arrays used in the compiler */
 int _PyCompile_EnsureArrayLargeEnough(
diff --git a/Lib/test/test_compile.py b/Lib/test/test_compile.py
index f22761f..736eff3 100644
--- a/Lib/test/test_compile.py
+++ b/Lib/test/test_compile.py
@@ -903,6 +903,28 @@ class TestSpecifics(unittest.TestCase):
                 self.assertIsNone(ns['with_const_expression'].__doc__)
 
     @support.cpython_only
+    def test_docstring_interactive_mode(self):
+        srcs = [
+            """def with_docstring():
+                "docstring"
+            """,
+            """class with_docstring:
+                "docstring"
+            """,
+        ]
+
+        for opt in [0, 1, 2]:
+            for src in srcs:
+                with self.subTest(opt=opt, src=src):
+                    code = compile(textwrap.dedent(src), "<test>", "single", optimize=opt)
+                    ns = {}
+                    exec(code, ns)
+                    if opt < 2:
+                        self.assertEqual(ns['with_docstring'].__doc__, "docstring")
+                    else:
+                        self.assertIsNone(ns['with_docstring'].__doc__)
+
+    @support.cpython_only
     def test_docstring_omitted(self):
         # See gh-115347
         src = textwrap.dedent("""
@@ -919,12 +941,13 @@ class TestSpecifics(unittest.TestCase):
                 return h
         """)
         for opt in [-1, 0, 1, 2]:
-            with self.subTest(opt=opt):
-                code = compile(src, "<test>", "exec", optimize=opt)
-                output = io.StringIO()
-                with contextlib.redirect_stdout(output):
-                    dis.dis(code)
-                self.assertNotIn('NOP' , output.getvalue())
+            for mode in ["exec", "single"]:
+                with self.subTest(opt=opt, mode=mode):
+                    code = compile(src, "<test>", mode, optimize=opt)
+                    output = io.StringIO()
+                    with contextlib.redirect_stdout(output):
+                        dis.dis(code)
+                    self.assertNotIn('NOP', output.getvalue())
 
     def test_dont_merge_constants(self):
         # Issue #25843: compile() must not merge constants which are equal
diff --git a/Misc/NEWS.d/next/Core_and_Builtins/2024-09-12-21-53-26.gh-issue-124022.fQzUiW.rst b/Misc/NEWS.d/next/Core_and_Builtins/2024-09-12-21-53-26.gh-issue-124022.fQzUiW.rst
new file mode 100644
index 0000000..90a77a5
--- /dev/null
+++ b/Misc/NEWS.d/next/Core_and_Builtins/2024-09-12-21-53-26.gh-issue-124022.fQzUiW.rst
@@ -0,0 +1 @@
+Fix bug where docstring is removed from classes in interactive mode.
diff --git a/Python/codegen.c b/Python/codegen.c
index 2ca5db1..5565d30 100644
--- a/Python/codegen.c
+++ b/Python/codegen.c
@@ -746,7 +746,7 @@ _PyCodegen_Expression(compiler *c, expr_ty e)
    and for annotations. */
 
 int
-_PyCodegen_Body(compiler *c, location loc, asdl_stmt_seq *stmts)
+_PyCodegen_Body(compiler *c, location loc, asdl_stmt_seq *stmts, bool is_interactive)
 {
     /* If from __future__ import annotations is active,
      * every annotated class and module should have __annotations__.
@@ -758,7 +758,7 @@ _PyCodegen_Body(compiler *c, location loc, asdl_stmt_seq *stmts)
         return SUCCESS;
     }
     Py_ssize_t first_instr = 0;
-    if (!IS_INTERACTIVE(c)) {
+    if (!is_interactive) { /* A string literal on REPL prompt is not a docstring */
         PyObject *docstring = _PyAST_GetDocString(stmts);
         if (docstring) {
             first_instr = 1;
@@ -1432,7 +1432,7 @@ codegen_class_body(compiler *c, stmt_ty s, int firstlineno)
         ADDOP_N_IN_SCOPE(c, loc, STORE_DEREF, &_Py_ID(__classdict__), cellvars);
     }
     /* compile the body proper */
-    RETURN_IF_ERROR_IN_SCOPE(c, _PyCodegen_Body(c, loc, s->v.ClassDef.body));
+    RETURN_IF_ERROR_IN_SCOPE(c, _PyCodegen_Body(c, loc, s->v.ClassDef.body, false));
     PyObject *static_attributes = _PyCompile_StaticAttributesAsTuple(c);
     if (static_attributes == NULL) {
         _PyCompile_ExitScope(c);
diff --git a/Python/compile.c b/Python/compile.c
index d54c320..e1d2c30 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -790,13 +790,13 @@ compiler_codegen(compiler *c, mod_ty mod)
     switch (mod->kind) {
     case Module_kind: {
         asdl_stmt_seq *stmts = mod->v.Module.body;
-        RETURN_IF_ERROR(_PyCodegen_Body(c, start_location(stmts), stmts));
+        RETURN_IF_ERROR(_PyCodegen_Body(c, start_location(stmts), stmts, false));
         break;
     }
     case Interactive_kind: {
         c->c_interactive = 1;
         asdl_stmt_seq *stmts = mod->v.Interactive.body;
-        RETURN_IF_ERROR(_PyCodegen_Body(c, start_location(stmts), stmts));
+        RETURN_IF_ERROR(_PyCodegen_Body(c, start_location(stmts), stmts, true));
         break;
     }
     case Expression_kind: {