python · pablogsal · Mar 2, 2020 · Dec 30, 2019 · Jan 2, 2020 · Mar 1, 2020
diff --git a/Lib/ast.py b/Lib/ast.py
@@ -667,6 +667,22 @@ def set_precedence(self, precedence, *nodes):
        for node in nodes:
            self._precedences[node] = precedence

+    def get_raw_docstring(self, node):
+        """If a docstring node is found in the body of the *node* parameter,
+        return that docstring node, None otherwise.
+
+        Logic mirrored from ``_PyAST_GetDocString``."""
+        if not isinstance(
+            node, (AsyncFunctionDef, FunctionDef, ClassDef, Module)
+        ) or len(node.body) < 1:
+            return None
+        node = node.body[0]
+        if not isinstance(node, Expr):
+            return None
+        node = node.value
+        if isinstance(node, Constant) and isinstance(node.value, str):
+            return node
+
    def traverse(self, node):
        if isinstance(node, list):
            for item in node:
@@ -681,9 +697,15 @@ def visit(self, node):
        self.traverse(node)
        return "".join(self._source)

+    def _write_docstring_and_traverse_body(self, node):
+        if (docstring := self.get_raw_docstring(node)):
+            self._write_docstring(docstring)
+            self.traverse(node.body[1:])
+        else:
+            self.traverse(node.body)
+
    def visit_Module(self, node):
-        for subnode in node.body:
-            self.traverse(subnode)
+        self._write_docstring_and_traverse_body(node)

    def visit_Expr(self, node):
        self.fill()
@@ -850,15 +872,15 @@ def visit_ClassDef(self, node):
                self.traverse(e)

        with self.block():
-            self.traverse(node.body)
+            self._write_docstring_and_traverse_body(node)

    def visit_FunctionDef(self, node):
-        self.__FunctionDef_helper(node, "def")
+        self._function_helper(node, "def")

    def visit_AsyncFunctionDef(self, node):
-        self.__FunctionDef_helper(node, "async def")
+        self._function_helper(node, "async def")

-    def __FunctionDef_helper(self, node, fill_suffix):
+    def _function_helper(self, node, fill_suffix):
        self.write("\n")
        for deco in node.decorator_list:
            self.fill("@")
@@ -871,15 +893,15 @@ def __FunctionDef_helper(self, node, fill_suffix):
            self.write(" -> ")
            self.traverse(node.returns)
        with self.block():
-            self.traverse(node.body)
+            self._write_docstring_and_traverse_body(node)

    def visit_For(self, node):
-        self.__For_helper("for ", node)
+        self._for_helper("for ", node)

    def visit_AsyncFor(self, node):
-        self.__For_helper("async for ", node)
+        self._for_helper("async for ", node)

-    def __For_helper(self, fill, node):
+    def _for_helper(self, fill, node):
        self.fill(fill)
        self.traverse(node.target)
        self.write(" in ")
@@ -974,6 +996,19 @@ def _fstring_FormattedValue(self, node, write):
    def visit_Name(self, node):
        self.write(node.id)

+    def _write_docstring(self, node):
+        self.fill()
+        if node.kind == "u":
+            self.write("u")
+
+        # Preserve quotes in the docstring by escaping them
+        value = node.value.replace("\\", "\\\\")
+        value = value.replace('"""', '""\"')
+        if value[-1] == '"':
+            value = value.replace('"', '\\"', -1)
+
+        self.write(f'"""{value}"""')
+
    def _write_constant(self, value):
        if isinstance(value, (float, complex)):
            # Substitute overflowing decimal literal for AST infinities.