do not add None to co_consts if no docstring

python · markshannon · Oct 30, 2024 · Oct 17, 2024 · Oct 20, 2024 · Oct 22, 2024
commit b6edafd8047aa8e5ad9978e4c33a312eae4ac7e0
diff --git a/Include/internal/pycore_symtable.h b/Include/internal/pycore_symtable.h
@@ -123,6 +123,7 @@ typedef struct _symtable_entry {
     unsigned ste_comp_iter_target : 1; /* true if visiting comprehension target */
     unsigned ste_can_see_class_scope : 1; /* true if this block can see names bound in an
                                              enclosing class scope */
+    unsigned ste_has_docstring : 1; /* true if docstring present */
     int ste_comp_iter_expr; /* non-zero if visiting a comprehension range expression */
     _Py_SourceLocation ste_loc; /* source location of block */
     struct _symtable_entry *ste_annotation_block; /* symbol table entry for this entry's annotations */

@@ -772,6 +772,10 @@ _PyCodegen_Body(compiler *c, location loc, asdl_stmt_seq *stmts, bool is_interac
             assert(st->kind == Expr_kind);
             location loc = LOC(st->v.Expr.value);
             ADDOP_LOAD_CONST(c, loc, cleandoc);
+
+            PySTEntryObject *ste = SYMTABLE_ENTRY(c);
+            ste->ste_has_docstring = 1;
+
             Py_DECREF(cleandoc);
             RETURN_IF_ERROR(codegen_nameop(c, NO_LOCATION, &_Py_ID(__doc__), Store));
         }
@@ -1225,21 +1229,22 @@ 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);
+    PySTEntryObject *ste = SYMTABLE_ENTRY(c);
     if (docstring) {
         first_instr = 1;
         docstring = _PyCompile_CleanDoc(docstring);
         if (docstring == NULL) {
             _PyCompile_ExitScope(c);
             return ERROR;
         }
+        Py_ssize_t idx = _PyCompile_AddConst(c, docstring);
+        ste->ste_has_docstring = 1;
+        RETURN_IF_ERROR_IN_SCOPE(c, idx < 0 ? ERROR : SUCCESS);
     }
-    Py_ssize_t idx = _PyCompile_AddConst(c, docstring ? docstring : Py_None);
     Py_XDECREF(docstring);
-    RETURN_IF_ERROR_IN_SCOPE(c, idx < 0 ? ERROR : SUCCESS);
 
     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 */

@@ -1285,6 +1285,8 @@ compute_code_flags(compiler *c)
             flags |= CO_VARARGS;
         if (ste->ste_varkeywords)
             flags |= CO_VARKEYWORDS;
+        if (ste->ste_has_docstring)
+            flags |= CO_HAS_DOCSTRING;
     }
 
     if (ste->ste_coroutine && !ste->ste_generator) {

@@ -136,6 +136,8 @@ ste_new(struct symtable *st, identifier name, _Py_block_ty block,
     ste->ste_needs_classdict = 0;
     ste->ste_annotation_block = NULL;
 
+    ste->ste_has_docstring = 0;
+
     ste->ste_symbols = PyDict_New();
     ste->ste_varnames = PyList_New(0);
     ste->ste_children = PyList_New(0);