set ste_has_docstring for class and module

Author: xuantengh

Python/codegen.c

@@ -1231,10 +1231,12 @@ 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) {
+        assert(ste->ste_has_docstring);
         first_instr = 1;
         docstring = _PyCompile_CleanDoc(docstring);
         if (docstring == NULL) {
@@ -1248,7 +1250,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 */

Python/symtable.c

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