This is the proper way to set the docstring

src/analysis/scoping_analysis.cpp

@@ -62,8 +62,17 @@ struct ScopingAnalysis::ScopeNameUsage {
 
     ScopeNameUsage(AST* node, ScopeNameUsage* parent) : node(node), parent(parent) {
         if (node->type == AST_TYPE::ClassDef) {
+            AST_ClassDef* classdef = ast_cast<AST_ClassDef>(node);
+
             // classes have an implicit write to "__module__"
             written.insert("__module__");
+
+            if (classdef->body.size() && classdef->body[0]->type == AST_TYPE::Expr) {
+                AST_Expr* first_expr = ast_cast<AST_Expr>(classdef->body[0]);
+                if (first_expr->value->type == AST_TYPE::Str) {
+                    written.insert("__doc__");
+                }
+            }
         }
     }
 };

src/core/cfg.cpp

@@ -1668,15 +1668,25 @@ CFG* computeCFG(SourceInfo* source, std::vector<AST_stmt*> body) {
     CFG* rtn = new CFG();
     CFGVisitor visitor(source->ast->type, rtn);
 
-    // In a classdef, the "__module__" attribute is immediately available:
     if (source->ast->type == AST_TYPE::ClassDef) {
+        // A classdef always starts with "__module__ = __name__"
         Box* module_name = source->parent_module->getattr("__name__", NULL, NULL);
         assert(module_name->cls == str_cls);
         AST_Assign* module_assign = new AST_Assign();
         module_assign->targets.push_back(makeName("__module__", AST_TYPE::Store));
         module_assign->value = new AST_Str(static_cast<BoxedString*>(module_name)->s);
-
         visitor.push_back(module_assign);
+
+        // If the first statement is just a single string, transform it to an assignment to __doc__
+        if (body.size() && body[0]->type == AST_TYPE::Expr) {
+            AST_Expr* first_expr = ast_cast<AST_Expr>(body[0]);
+            if (first_expr->value->type == AST_TYPE::Str) {
+                AST_Assign* doc_assign = new AST_Assign();
+                doc_assign->targets.push_back(makeName("__doc__", AST_TYPE::Store));
+                doc_assign->value = first_expr->value;
+                visitor.push_back(doc_assign);
+            }
+        }
     }
 
     for (int i = 0; i < body.size(); i++) {
@@ -1691,17 +1701,6 @@ CFG* computeCFG(SourceInfo* source, std::vector<AST_stmt*> body) {
         auto written_names = scope_info->getClassDefLocalNames();
         AST_Dict* rtn_dict = new AST_Dict();
 
-        // It'd be ok to add __doc__ to the dict multiple times, since the last one would win
-        if (written_names.count("__doc__") == 0) {
-            if (body.size() && body[0]->type == AST_TYPE::Expr) {
-                AST_Expr* first_expr = ast_cast<AST_Expr>(body[0]);
-                if (first_expr->value->type == AST_TYPE::Str) {
-                    rtn_dict->keys.push_back(new AST_Str("__doc__"));
-                    rtn_dict->values.push_back(first_expr->value);
-                }
-            }
-        }
-
         // Even if the user never explicitly wrote to __module__, there was an
         // implicit write:
         assert(written_names.count("__module__"));

test/tests/docstrings.py

@@ -11,6 +11,7 @@ print C1.__doc__
 class C2(object):
     "doc1"
     "doc2"
+    print __doc__
 print C2.__doc__
 
 class C3(object):

tools/tester.py

@@ -189,7 +189,7 @@ def run_test(fn, check_stats, run_memcheck):
             out_fd, out_fn = tempfile.mkstemp()
             os.fdopen(exp_fd, 'w').write(expected_out)
             os.fdopen(out_fd, 'w').write(out)
-            p = subprocess.Popen(["diff", "-a", exp_fn, out_fn], stdout=subprocess.PIPE, preexec_fn=set_ulimits)
+            p = subprocess.Popen(["diff", "-C2", "-a", exp_fn, out_fn], stdout=subprocess.PIPE, preexec_fn=set_ulimits)
             diff = p.stdout.read()
             assert p.wait() in (0, 1)
             raise Exception("Failed on %s:\n%s" % (fn, diff))