bpo-29463: Add docstring field to some AST nodes. (#46)

* bpo-29463: Add docstring field to some AST nodes.

ClassDef, ModuleDef, FunctionDef, and AsyncFunctionDef has docstring
field for now.  It was first statement of there body.

* fix document.  thanks travis!

* doc fixes

1 files changed, 63 insertions, 26 deletions

diff --git a/Python/ast.c b/Python/ast.c
index 5c5738f..55beef5 100644
--- a/Python/ast.c
+++ b/Python/ast.c
@@ -366,9 +366,12 @@ validate_assignlist(asdl_seq *targets, expr_context_ty ctx)
 }
 
 static int
-validate_body(asdl_seq *body, const char *owner)
+validate_body(asdl_seq *body, const char *owner, int allowempty)
 {
-    return validate_nonempty_seq(body, "body", owner) && validate_stmts(body);
+    if (!allowempty && !validate_nonempty_seq(body, "body", owner)) {
+        return 0;
+    }
+    return validate_stmts(body);
 }
 
 static int
@@ -377,13 +380,15 @@ validate_stmt(stmt_ty stmt)
     int i;
     switch (stmt->kind) {
     case FunctionDef_kind:
-        return validate_body(stmt->v.FunctionDef.body, "FunctionDef") &&
+        return validate_body(stmt->v.FunctionDef.body, "FunctionDef",
+                             stmt->v.FunctionDef.docstring != NULL) &&
             validate_arguments(stmt->v.FunctionDef.args) &&
             validate_exprs(stmt->v.FunctionDef.decorator_list, Load, 0) &&
             (!stmt->v.FunctionDef.returns ||
              validate_expr(stmt->v.FunctionDef.returns, Load));
     case ClassDef_kind:
-        return validate_body(stmt->v.ClassDef.body, "ClassDef") &&
+        return validate_body(stmt->v.ClassDef.body, "ClassDef",
+                             stmt->v.ClassDef.docstring != NULL) &&
             validate_exprs(stmt->v.ClassDef.bases, Load, 0) &&
             validate_keywords(stmt->v.ClassDef.keywords) &&
             validate_exprs(stmt->v.ClassDef.decorator_list, Load, 0);
@@ -411,20 +416,20 @@ validate_stmt(stmt_ty stmt)
     case For_kind:
         return validate_expr(stmt->v.For.target, Store) &&
             validate_expr(stmt->v.For.iter, Load) &&
-            validate_body(stmt->v.For.body, "For") &&
+            validate_body(stmt->v.For.body, "For", 0) &&
             validate_stmts(stmt->v.For.orelse);
     case AsyncFor_kind:
         return validate_expr(stmt->v.AsyncFor.target, Store) &&
             validate_expr(stmt->v.AsyncFor.iter, Load) &&
-            validate_body(stmt->v.AsyncFor.body, "AsyncFor") &&
+            validate_body(stmt->v.AsyncFor.body, "AsyncFor", 0) &&
             validate_stmts(stmt->v.AsyncFor.orelse);
     case While_kind:
         return validate_expr(stmt->v.While.test, Load) &&
-            validate_body(stmt->v.While.body, "While") &&
+            validate_body(stmt->v.While.body, "While", 0) &&
             validate_stmts(stmt->v.While.orelse);
     case If_kind:
         return validate_expr(stmt->v.If.test, Load) &&
-            validate_body(stmt->v.If.body, "If") &&
+            validate_body(stmt->v.If.body, "If", 0) &&
             validate_stmts(stmt->v.If.orelse);
     case With_kind:
         if (!validate_nonempty_seq(stmt->v.With.items, "items", "With"))
@@ -435,7 +440,7 @@ validate_stmt(stmt_ty stmt)
                 (item->optional_vars && !validate_expr(item->optional_vars, Store)))
                 return 0;
         }
-        return validate_body(stmt->v.With.body, "With");
+        return validate_body(stmt->v.With.body, "With", 0);
     case AsyncWith_kind:
         if (!validate_nonempty_seq(stmt->v.AsyncWith.items, "items", "AsyncWith"))
             return 0;
@@ -445,7 +450,7 @@ validate_stmt(stmt_ty stmt)
                 (item->optional_vars && !validate_expr(item->optional_vars, Store)))
                 return 0;
         }
-        return validate_body(stmt->v.AsyncWith.body, "AsyncWith");
+        return validate_body(stmt->v.AsyncWith.body, "AsyncWith", 0);
     case Raise_kind:
         if (stmt->v.Raise.exc) {
             return validate_expr(stmt->v.Raise.exc, Load) &&
@@ -457,7 +462,7 @@ validate_stmt(stmt_ty stmt)
         }
         return 1;
     case Try_kind:
-        if (!validate_body(stmt->v.Try.body, "Try"))
+        if (!validate_body(stmt->v.Try.body, "Try", 0))
             return 0;
         if (!asdl_seq_LEN(stmt->v.Try.handlers) &&
             !asdl_seq_LEN(stmt->v.Try.finalbody)) {
@@ -473,7 +478,7 @@ validate_stmt(stmt_ty stmt)
             excepthandler_ty handler = asdl_seq_GET(stmt->v.Try.handlers, i);
             if ((handler->v.ExceptHandler.type &&
                  !validate_expr(handler->v.ExceptHandler.type, Load)) ||
-                !validate_body(handler->v.ExceptHandler.body, "ExceptHandler"))
+                !validate_body(handler->v.ExceptHandler.body, "ExceptHandler", 0))
                 return 0;
         }
         return (!asdl_seq_LEN(stmt->v.Try.finalbody) ||
@@ -498,7 +503,8 @@ validate_stmt(stmt_ty stmt)
     case Expr_kind:
         return validate_expr(stmt->v.Expr.value, Load);
     case AsyncFunctionDef_kind:
-        return validate_body(stmt->v.AsyncFunctionDef.body, "AsyncFunctionDef") &&
+        return validate_body(stmt->v.AsyncFunctionDef.body, "AsyncFunctionDef",
+                             stmt->v.AsyncFunctionDef.docstring != NULL) &&
             validate_arguments(stmt->v.AsyncFunctionDef.args) &&
             validate_exprs(stmt->v.AsyncFunctionDef.decorator_list, Load, 0) &&
             (!stmt->v.AsyncFunctionDef.returns ||
@@ -594,7 +600,9 @@ struct compiling {
 static asdl_seq *seq_for_testlist(struct compiling *, const node *);
 static expr_ty ast_for_expr(struct compiling *, const node *);
 static stmt_ty ast_for_stmt(struct compiling *, const node *);
-static asdl_seq *ast_for_suite(struct compiling *, const node *);
+static asdl_seq *ast_for_body(struct compiling *c, const node *n,
+                              string *docstring);
+static string docstring_from_stmts(asdl_seq *stmts);
 static asdl_seq *ast_for_exprlist(struct compiling *, const node *,
                                   expr_context_ty);
 static expr_ty ast_for_testlist(struct compiling *, const node *);
@@ -806,7 +814,7 @@ PyAST_FromNodeObject(const node *n, PyCompilerFlags *flags,
                     }
                 }
             }
-            res = Module(stmts, arena);
+            res = Module(stmts, docstring_from_stmts(stmts), arena);
             break;
         case eval_input: {
             expr_ty testlist_ast;
@@ -1593,6 +1601,7 @@ ast_for_funcdef_impl(struct compiling *c, const node *n,
     arguments_ty args;
     asdl_seq *body;
     expr_ty returns = NULL;
+    string docstring;
     int name_i = 1;
 
     REQ(n, funcdef);
@@ -1611,17 +1620,17 @@ ast_for_funcdef_impl(struct compiling *c, const node *n,
             return NULL;
         name_i += 2;
     }
-    body = ast_for_suite(c, CHILD(n, name_i + 3));
+    body = ast_for_body(c, CHILD(n, name_i + 3), &docstring);
     if (!body)
         return NULL;
 
     if (is_async)
         return AsyncFunctionDef(name, args, body, decorator_seq, returns,
-                                LINENO(n),
+                                docstring, LINENO(n),
                                 n->n_col_offset, c->c_arena);
     else
         return FunctionDef(name, args, body, decorator_seq, returns,
-                           LINENO(n),
+                           docstring, LINENO(n),
                            n->n_col_offset, c->c_arena);
 }
 
@@ -3514,6 +3523,32 @@ ast_for_suite(struct compiling *c, const node *n)
     return seq;
 }
 
+static string
+docstring_from_stmts(asdl_seq *stmts)
+{
+    if (stmts && stmts->size) {
+        stmt_ty s = (stmt_ty)asdl_seq_GET(stmts, 0);
+        /* If first statement is a literal string, it's the doc string. */
+        if (s->kind == Expr_kind && s->v.Expr.value->kind == Str_kind) {
+            string doc = s->v.Expr.value->v.Str.s;
+            /* not very efficient, but simple */
+            memmove(&asdl_seq_GET(stmts, 0), &asdl_seq_GET(stmts, 1),
+                    (stmts->size - 1) * sizeof(void*));
+            stmts->size--;
+            return doc;
+        }
+    }
+    return NULL;
+}
+
+static asdl_seq *
+ast_for_body(struct compiling *c, const node *n, string *docstring)
+{
+    asdl_seq *stmts = ast_for_suite(c, n);
+    *docstring = docstring_from_stmts(stmts);
+    return stmts;
+}
+
 static stmt_ty
 ast_for_if_stmt(struct compiling *c, const node *n)
 {
@@ -3898,12 +3933,13 @@ ast_for_classdef(struct compiling *c, const node *n, asdl_seq *decorator_seq)
     /* classdef: 'class' NAME ['(' arglist ')'] ':' suite */
     PyObject *classname;
     asdl_seq *s;
+    string docstring;
     expr_ty call;
 
     REQ(n, classdef);
 
     if (NCH(n) == 4) { /* class NAME ':' suite */
-        s = ast_for_suite(c, CHILD(n, 3));
+        s = ast_for_body(c, CHILD(n, 3), &docstring);
         if (!s)
             return NULL;
         classname = NEW_IDENTIFIER(CHILD(n, 1));
@@ -3911,12 +3947,12 @@ ast_for_classdef(struct compiling *c, const node *n, asdl_seq *decorator_seq)
             return NULL;
         if (forbidden_name(c, classname, CHILD(n, 3), 0))
             return NULL;
-        return ClassDef(classname, NULL, NULL, s, decorator_seq, LINENO(n),
-                        n->n_col_offset, c->c_arena);
+        return ClassDef(classname, NULL, NULL, s, decorator_seq, docstring,
+                        LINENO(n), n->n_col_offset, c->c_arena);
     }
 
     if (TYPE(CHILD(n, 3)) == RPAR) { /* class NAME '(' ')' ':' suite */
-        s = ast_for_suite(c, CHILD(n,5));
+        s = ast_for_body(c, CHILD(n, 5), &docstring);
         if (!s)
             return NULL;
         classname = NEW_IDENTIFIER(CHILD(n, 1));
@@ -3924,8 +3960,8 @@ ast_for_classdef(struct compiling *c, const node *n, asdl_seq *decorator_seq)
             return NULL;
         if (forbidden_name(c, classname, CHILD(n, 3), 0))
             return NULL;
-        return ClassDef(classname, NULL, NULL, s, decorator_seq, LINENO(n),
-                        n->n_col_offset, c->c_arena);
+        return ClassDef(classname, NULL, NULL, s, decorator_seq, docstring,
+                        LINENO(n), n->n_col_offset, c->c_arena);
     }
 
     /* class NAME '(' arglist ')' ':' suite */
@@ -3941,7 +3977,7 @@ ast_for_classdef(struct compiling *c, const node *n, asdl_seq *decorator_seq)
         if (!call)
             return NULL;
     }
-    s = ast_for_suite(c, CHILD(n, 6));
+    s = ast_for_body(c, CHILD(n, 6), &docstring);
     if (!s)
         return NULL;
     classname = NEW_IDENTIFIER(CHILD(n, 1));
@@ -3951,7 +3987,8 @@ ast_for_classdef(struct compiling *c, const node *n, asdl_seq *decorator_seq)
         return NULL;
 
     return ClassDef(classname, call->v.Call.args, call->v.Call.keywords, s,
-                    decorator_seq, LINENO(n), n->n_col_offset, c->c_arena);
+                    decorator_seq, docstring, LINENO(n), n->n_col_offset,
+                    c->c_arena);
 }
 
 static stmt_ty