calculate ste_has_docstring in symtable.c

xuantengh · xuantengh · commit 13602e2b6647 · 2024-10-29T14:06:20.000+08:00
diff --git a/Python/codegen.c b/Python/codegen.c
@@ -773,9 +773,6 @@ _PyCodegen_Body(compiler *c, location loc, asdl_stmt_seq *stmts, bool is_interac
             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));
         }
@@ -1229,7 +1226,6 @@ 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);
@@ -1238,13 +1234,13 @@ codegen_function_body(compiler *c, stmt_ty s, int is_async, Py_ssize_t funcflags
             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_XDECREF(docstring);
 
     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 */
diff --git a/Python/symtable.c b/Python/symtable.c
@@ -1843,6 +1843,10 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
             return 0;
         }
 
+        if (_PyAST_GetDocString(s->v.FunctionDef.body)) {
+            new_ste->ste_has_docstring = 1;
+        }
+
         if (!symtable_visit_annotations(st, s, s->v.FunctionDef.args,
                                         s->v.FunctionDef.returns, new_ste)) {
             Py_DECREF(new_ste);
@@ -2170,6 +2174,10 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
             return 0;
         }
 
+        if (_PyAST_GetDocString(s->v.FunctionDef.body)) {
+            new_ste->ste_has_docstring = 1;
+        }
+
         if (!symtable_visit_annotations(st, s, s->v.AsyncFunctionDef.args,
                                         s->v.AsyncFunctionDef.returns, new_ste)) {
             Py_DECREF(new_ste);
