bpo-33691: Add _PyAST_GetDocString(). (GH-7236)

4 files changed, 43 insertions, 58 deletions

diff --git a/Python/ast.c b/Python/ast.c
index 43bd786015..3b4cd162fc 100644
--- a/Python/ast.c
+++ b/Python/ast.c
@@ -5260,3 +5260,23 @@ error:
     FstringParser_Dealloc(&state);
     return NULL;
 }
+
+PyObject *
+_PyAST_GetDocString(asdl_seq *body)
+{
+    if (!asdl_seq_LEN(body)) {
+        return NULL;
+    }
+    stmt_ty st = (stmt_ty)asdl_seq_GET(body, 0);
+    if (st->kind != Expr_kind) {
+        return NULL;
+    }
+    expr_ty e = st->v.Expr.value;
+    if (e->kind == Str_kind) {
+        return e->v.Str.s;
+    }
+    if (e->kind == Constant_kind && PyUnicode_CheckExact(e->v.Constant.value)) {
+        return e->v.Constant.value;
+    }
+    return NULL;
+}
diff --git a/Python/ast_opt.c b/Python/ast_opt.c
index a998d1fb60..5e57638e3d 100644
--- a/Python/ast_opt.c
+++ b/Python/ast_opt.c
@@ -1,6 +1,8 @@
 /* AST Optimizer */
 #include "Python.h"
 #include "Python-ast.h"
+#include "node.h"
+#include "ast.h"
 
 
 /* TODO: is_const and get_const_value are copied from Python/compile.c.
@@ -468,36 +470,18 @@ static int astfold_excepthandler(excepthandler_ty node_, PyArena *ctx_, int opti
 }
 
 static int
-isdocstring(stmt_ty s)
-{
-    if (s->kind != Expr_kind)
-        return 0;
-    if (s->v.Expr.value->kind == Str_kind)
-        return 1;
-    if (s->v.Expr.value->kind == Constant_kind)
-        return PyUnicode_CheckExact(s->v.Expr.value->v.Constant.value);
-    return 0;
-}
-
-static int
 astfold_body(asdl_seq *stmts, PyArena *ctx_, int optimize_)
 {
-    if (!asdl_seq_LEN(stmts)) {
-        return 1;
-    }
-    int docstring = isdocstring((stmt_ty)asdl_seq_GET(stmts, 0));
+    int docstring = _PyAST_GetDocString(stmts) != NULL;
     CALL_SEQ(astfold_stmt, stmt_ty, stmts);
-    if (docstring) {
-        return 1;
-    }
-    stmt_ty st = (stmt_ty)asdl_seq_GET(stmts, 0);
-    if (isdocstring(st)) {
+    if (!docstring && _PyAST_GetDocString(stmts) != NULL) {
+        stmt_ty st = (stmt_ty)asdl_seq_GET(stmts, 0);
         asdl_seq *values = _Py_asdl_seq_new(1, ctx_);
         if (!values) {
             return 0;
         }
         asdl_seq_SET(values, 0, st->v.Expr.value);
-        expr_ty expr = _Py_JoinedStr(values, st->lineno, st->col_offset, ctx_);
+        expr_ty expr = JoinedStr(values, st->lineno, st->col_offset, ctx_);
         if (!expr) {
             return 0;
         }
diff --git a/Python/compile.c b/Python/compile.c
index 42ae508259..3528670ef6 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -1393,18 +1393,6 @@ compiler_addop_j(struct compiler *c, int opcode, basicblock *b, int absolute)
 }
 
 static int
-compiler_isdocstring(stmt_ty s)
-{
-    if (s->kind != Expr_kind)
-        return 0;
-    if (s->v.Expr.value->kind == Str_kind)
-        return 1;
-    if (s->v.Expr.value->kind == Constant_kind)
-        return PyUnicode_CheckExact(s->v.Expr.value->v.Constant.value);
-    return 0;
-}
-
-static int
 is_const(expr_ty e)
 {
     switch (e->kind) {
@@ -1603,6 +1591,7 @@ compiler_body(struct compiler *c, asdl_seq *stmts)
 {
     int i = 0;
     stmt_ty st;
+    PyObject *docstring;
 
     /* Set current line number to the line number of first statement.
        This way line number for SETUP_ANNOTATIONS will always
@@ -1619,14 +1608,17 @@ compiler_body(struct compiler *c, asdl_seq *stmts)
     }
     if (!asdl_seq_LEN(stmts))
         return 1;
-    st = (stmt_ty)asdl_seq_GET(stmts, 0);
     /* if not -OO mode, set docstring */
-    if (compiler_isdocstring(st) && c->c_optimize < 2) {
-        /* don't generate docstrings if -OO */
-        i = 1;
-        VISIT(c, expr, st->v.Expr.value);
-        if (!compiler_nameop(c, __doc__, Store))
-            return 0;
+    if (c->c_optimize < 2) {
+        docstring = _PyAST_GetDocString(stmts);
+        if (docstring) {
+            i = 1;
+            st = (stmt_ty)asdl_seq_GET(stmts, 0);
+            assert(st->kind == Expr_kind);
+            VISIT(c, expr, st->v.Expr.value);
+            if (!compiler_nameop(c, __doc__, Store))
+                return 0;
+        }
     }
     for (; i < asdl_seq_LEN(stmts); i++)
         VISIT(c, stmt, (stmt_ty)asdl_seq_GET(stmts, i));
@@ -1979,15 +1971,13 @@ static int
 compiler_function(struct compiler *c, stmt_ty s, int is_async)
 {
     PyCodeObject *co;
-    PyObject *qualname, *first_const = Py_None;
+    PyObject *qualname, *docstring = NULL;
     arguments_ty args;
     expr_ty returns;
     identifier name;
     asdl_seq* decos;
     asdl_seq *body;
-    stmt_ty st;
     Py_ssize_t i, funcflags;
-    int docstring;
     int annotations;
     int scope_type;
 
@@ -2034,15 +2024,10 @@ compiler_function(struct compiler *c, stmt_ty s, int is_async)
     }
 
     /* if not -OO mode, add docstring */
-    st = (stmt_ty)asdl_seq_GET(body, 0);
-    docstring = compiler_isdocstring(st);
-    if (docstring && c->c_optimize < 2) {
-        if (st->v.Expr.value->kind == Constant_kind)
-            first_const = st->v.Expr.value->v.Constant.value;
-        else
-            first_const = st->v.Expr.value->v.Str.s;
+    if (c->c_optimize < 2) {
+        docstring = _PyAST_GetDocString(body);
     }
-    if (compiler_add_const(c, first_const) < 0) {
+    if (compiler_add_const(c, docstring ? docstring : Py_None) < 0) {
         compiler_exit_scope(c);
         return 0;
     }
diff --git a/Python/future.c b/Python/future.c
index 03a97c865a..4ea6827723 100644
--- a/Python/future.c
+++ b/Python/future.c
@@ -5,6 +5,7 @@
 #include "graminit.h"
 #include "code.h"
 #include "symtable.h"
+#include "ast.h"
 
 #define UNDEFINED_FUTURE_FEATURE "future feature %.100s is not defined"
 #define ERR_LATE_FUTURE \
@@ -63,7 +64,6 @@ static int
 future_parse(PyFutureFeatures *ff, mod_ty mod, PyObject *filename)
 {
     int i, done = 0, prev_line = 0;
-    stmt_ty first;
 
     if (!(mod->kind == Module_kind || mod->kind == Interactive_kind))
         return 1;
@@ -80,11 +80,7 @@ future_parse(PyFutureFeatures *ff, mod_ty mod, PyObject *filename)
     */
 
     i = 0;
-    first = (stmt_ty)asdl_seq_GET(mod->v.Module.body, i);
-    if (first->kind == Expr_kind
-        && (first->v.Expr.value->kind == Str_kind
-            || (first->v.Expr.value->kind == Constant_kind
-                && PyUnicode_CheckExact(first->v.Expr.value->v.Constant.value))))
+    if (_PyAST_GetDocString(mod->v.Module.body) != NULL)
         i++;
 
     for (; i < asdl_seq_LEN(mod->v.Module.body); i++) {