gh-126072: Set docstring attribute for module and class (#126231)

Author: xuantengh

Lib/test/test_code.py

@@ -178,6 +178,20 @@
 nlocals: 3
 flags: 3
 consts: ("'hello'", "'world'")
+
+>>> class class_with_docstring:
+...     '''This is a docstring for class'''
+...     '''This line is not docstring'''
+...     pass
+
+>>> print(class_with_docstring.__doc__)
+This is a docstring for class
+
+>>> class class_without_docstring:
+...     pass
+
+>>> print(class_without_docstring.__doc__)
+None
 """
 
 import copy
@@ -854,6 +868,33 @@ def f():
             3 * [(42, 42, None, None)],
         )
 
+    @cpython_only
+    def test_docstring_under_o2(self):
+        code = textwrap.dedent('''
+            def has_docstring(x, y):
+                """This is a first-line doc string"""
+                """This is a second-line doc string"""
+                a = x + y
+                b = x - y
+                return a, b
+
+
+            def no_docstring(x):
+                def g(y):
+                    return x + y
+                return g
+
+
+            async def async_func():
+                """asynf function doc string"""
+                pass
+
+
+            for func in [has_docstring, no_docstring(4), async_func]:
+                assert(func.__doc__ is None)
+            ''')
+
+        rc, out, err = assert_python_ok('-OO', '-c', code)
 
 if check_impl_detail(cpython=True) and ctypes is not None:
     py = ctypes.pythonapi

Lib/test/test_compile.py

@@ -342,6 +342,10 @@ def test_lambda_doc(self):
         l = lambda: "foo"
         self.assertIsNone(l.__doc__)
 
+    def test_lambda_consts(self):
+        l = lambda: "this is the only const"
+        self.assertEqual(l.__code__.co_consts, ("this is the only const",))
+
     def test_encoding(self):
         code = b'# -*- coding: badencoding -*-\npass\n'
         self.assertRaises(SyntaxError, compile, code, 'tmp', 'exec')
@@ -790,10 +794,10 @@ def check_same_constant(const):
         # Merge constants in tuple or frozenset
         f1, f2 = lambda: "not a name", lambda: ("not a name",)
         f3 = lambda x: x in {("not a name",)}
-        self.assertIs(f1.__code__.co_consts[1],
-                      f2.__code__.co_consts[1][0])
-        self.assertIs(next(iter(f3.__code__.co_consts[1])),
-                      f2.__code__.co_consts[1])
+        self.assertIs(f1.__code__.co_consts[0],
+                      f2.__code__.co_consts[0][0])
+        self.assertIs(next(iter(f3.__code__.co_consts[0])),
+                      f2.__code__.co_consts[0])
 
         # {0} is converted to a constant frozenset({0}) by the peephole
         # optimizer
@@ -902,6 +906,9 @@ def with_fstring():
 
             def with_const_expression():
                 "also" + " not docstring"
+
+            def multiple_const_strings():
+                "not docstring " * 3
             """)
 
         for opt in [0, 1, 2]:
@@ -918,6 +925,7 @@ def with_const_expression():
                     self.assertIsNone(ns['two_strings'].__doc__)
                 self.assertIsNone(ns['with_fstring'].__doc__)
                 self.assertIsNone(ns['with_const_expression'].__doc__)
+                self.assertIsNone(ns['multiple_const_strings'].__doc__)
 
     @support.cpython_only
     def test_docstring_interactive_mode(self):

Misc/ACKS

@@ -820,6 +820,7 @@ Tomáš Hrnčiar
 Miro Hrončok
 Chiu-Hsiang Hsu
 Chih-Hao Huang
+Xuanteng Huang
 Christian Hudon
 Benoît Hudson
 Lawrence Hudson

Misc/NEWS.d/next/Core_and_Builtins/2024-10-31-21-49-00.gh-issue-126072.o9k8Ns.rst

@@ -0,0 +1,2 @@
+Following :gh:`126101`, for :ref:`codeobjects` like lambda, annotation and type alias,
+we no longer  add ``None`` to its :attr:`~codeobject.co_consts`.

Python/codegen.c

@@ -672,9 +672,7 @@ codegen_setup_annotations_scope(compiler *c, location loc,
         codegen_enter_scope(c, name, COMPILE_SCOPE_ANNOTATIONS,
                             key, loc.lineno, NULL, &umd));
 
-    // Insert None into consts to prevent an annotation
-    // appearing to be a docstring
-    _PyCompile_AddConst(c, Py_None);
+    assert(!SYMTABLE_ENTRY(c)->ste_has_docstring);
     // if .format != 1: raise NotImplementedError
     _Py_DECLARE_STR(format, ".format");
     ADDOP_I(c, loc, LOAD_FAST, 0);
@@ -770,16 +768,18 @@ _PyCodegen_Body(compiler *c, location loc, asdl_stmt_seq *stmts, bool is_interac
     /* If from __future__ import annotations is active,
      * every annotated class and module should have __annotations__.
      * Else __annotate__ is created when necessary. */
-    if ((FUTURE_FEATURES(c) & CO_FUTURE_ANNOTATIONS) && SYMTABLE_ENTRY(c)->ste_annotations_used) {
+    PySTEntryObject *ste = SYMTABLE_ENTRY(c);
+    if ((FUTURE_FEATURES(c) & CO_FUTURE_ANNOTATIONS) && ste->ste_annotations_used) {
         ADDOP(c, loc, SETUP_ANNOTATIONS);
     }
     if (!asdl_seq_LEN(stmts)) {
         return SUCCESS;
     }
     Py_ssize_t first_instr = 0;
     if (!is_interactive) { /* A string literal on REPL prompt is not a docstring */
-        PyObject *docstring = _PyAST_GetDocString(stmts);
-        if (docstring) {
+        if (ste->ste_has_docstring) {
+            PyObject *docstring = _PyAST_GetDocString(stmts);
+            assert(docstring);
             first_instr = 1;
             /* set docstring */
             assert(OPTIMIZATION_LEVEL(c) < 2);
@@ -1241,10 +1241,11 @@ codegen_function_body(compiler *c, stmt_ty s, int is_async, Py_ssize_t funcflags
     RETURN_IF_ERROR(
         codegen_enter_scope(c, name, scope_type, (void *)s, firstlineno, NULL, &umd));
 
+    PySTEntryObject *ste = SYMTABLE_ENTRY(c);
     Py_ssize_t first_instr = 0;
-    PyObject *docstring = _PyAST_GetDocString(body);
-    assert(OPTIMIZATION_LEVEL(c) < 2 || docstring == NULL);
-    if (docstring) {
+    if (ste->ste_has_docstring) {
+        PyObject *docstring = _PyAST_GetDocString(body);
+        assert(docstring);
         first_instr = 1;
         docstring = _PyCompile_CleanDoc(docstring);
         if (docstring == NULL) {
@@ -1258,7 +1259,6 @@ codegen_function_body(compiler *c, stmt_ty s, int is_async, Py_ssize_t funcflags
 
     NEW_JUMP_TARGET_LABEL(c, start);
     USE_LABEL(c, start);
-    PySTEntryObject *ste = SYMTABLE_ENTRY(c);
     bool add_stopiteration_handler = ste->ste_coroutine || ste->ste_generator;
     if (add_stopiteration_handler) {
         /* codegen_wrap_in_stopiteration_handler will push a block, so we need to account for that */
@@ -1600,9 +1600,8 @@ codegen_typealias_body(compiler *c, stmt_ty s)
     ADDOP_LOAD_CONST_NEW(c, loc, defaults);
     RETURN_IF_ERROR(
         codegen_setup_annotations_scope(c, LOC(s), s, name));
-    /* Make None the first constant, so the evaluate function can't have a
-        docstring. */
-    RETURN_IF_ERROR(_PyCompile_AddConst(c, Py_None));
+
+    assert(!SYMTABLE_ENTRY(c)->ste_has_docstring);
     VISIT_IN_SCOPE(c, expr, s->v.TypeAlias.value);
     ADDOP_IN_SCOPE(c, loc, RETURN_VALUE);
     PyCodeObject *co = _PyCompile_OptimizeAndAssemble(c, 0);
@@ -1898,9 +1897,7 @@ codegen_lambda(compiler *c, expr_ty e)
         codegen_enter_scope(c, &_Py_STR(anon_lambda), COMPILE_SCOPE_LAMBDA,
                             (void *)e, e->lineno, NULL, &umd));
 
-    /* Make None the first constant, so the lambda can't have a
-       docstring. */
-    RETURN_IF_ERROR(_PyCompile_AddConst(c, Py_None));
+    assert(!SYMTABLE_ENTRY(c)->ste_has_docstring);
 
     VISIT_IN_SCOPE(c, expr, e->v.Lambda.body);
     if (SYMTABLE_ENTRY(c)->ste_generator) {

Python/symtable.c

@@ -434,6 +434,9 @@ _PySymtable_Build(mod_ty mod, PyObject *filename, _PyFutureFeatures *future)
     switch (mod->kind) {
     case Module_kind:
         seq = mod->v.Module.body;
+        if (_PyAST_GetDocString(seq)) {
+            st->st_cur->ste_has_docstring = 1;
+        }
         for (i = 0; i < asdl_seq_LEN(seq); i++)
             if (!symtable_visit_stmt(st,
                         (stmt_ty)asdl_seq_GET(seq, i)))
@@ -1909,6 +1912,11 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
                 return 0;
             }
         }
+
+        if (_PyAST_GetDocString(s->v.ClassDef.body)) {
+            st->st_cur->ste_has_docstring = 1;
+        }
+
         VISIT_SEQ(st, stmt, s->v.ClassDef.body);
         if (!symtable_exit_block(st))
             return 0;