[Python-checkins] bpo-33851: Fix ast.get_docstring() for a node that lacks a docstring. (GH-7682)

[Serhiy Storchaka]

https://github.com/python/cpython/commit/08f127a3cad8ce4eb281d30d9488c91b0fd7cfed
commit: 08f127a3cad8ce4eb281d30d9488c91b0fd7cfed
branch: master
author: Serhiy Storchaka <storchaka at gmail.com>
committer: GitHub <noreply at github.com>
date: 2018-06-15T11:05:15+03:00
summary:

bpo-33851: Fix ast.get_docstring() for a node that lacks a docstring. (GH-7682)

files:
A Misc/NEWS.d/next/Library/2018-06-13-15-12-25.bpo-33851.SVbqlz.rst
M Lib/ast.py
M Lib/test/test_ast.py

diff --git a/Lib/ast.py b/Lib/ast.py
index 134d9d27582e..bfe346bba8e3 100644
--- a/Lib/ast.py
+++ b/Lib/ast.py
@@ -206,7 +206,7 @@ def get_docstring(node, clean=True):
     """
     if not isinstance(node, (AsyncFunctionDef, FunctionDef, ClassDef, Module)):
         raise TypeError("%r can't have docstrings" % node.__class__.__name__)
-    if not node.body:
+    if not(node.body and isinstance(node.body[0], Expr)):
         return None
     node = node.body[0].value
     if isinstance(node, Str):
@@ -215,7 +215,7 @@ def get_docstring(node, clean=True):
         text = node.value
     else:
         return None
-    if clean and text:
+    if clean:
         import inspect
         text = inspect.cleandoc(text)
     return text
diff --git a/Lib/test/test_ast.py b/Lib/test/test_ast.py
index ab32d9d389e6..7db40e7797f7 100644
--- a/Lib/test/test_ast.py
+++ b/Lib/test/test_ast.py
@@ -521,13 +521,44 @@ def test_iter_child_nodes(self):
         )
 
     def test_get_docstring(self):
+        node = ast.parse('"""line one\n  line two"""')
+        self.assertEqual(ast.get_docstring(node),
+                         'line one\nline two')
+
+        node = ast.parse('class foo:\n  """line one\n  line two"""')
+        self.assertEqual(ast.get_docstring(node.body[0]),
+                         'line one\nline two')
+
         node = ast.parse('def foo():\n  """line one\n  line two"""')
         self.assertEqual(ast.get_docstring(node.body[0]),
                          'line one\nline two')
 
         node = ast.parse('async def foo():\n  """spam\n  ham"""')
         self.assertEqual(ast.get_docstring(node.body[0]), 'spam\nham')
+
+    def test_get_docstring_none(self):
         self.assertIsNone(ast.get_docstring(ast.parse('')))
+        node = ast.parse('x = "not docstring"')
+        self.assertIsNone(ast.get_docstring(node))
+        node = ast.parse('def foo():\n  pass')
+        self.assertIsNone(ast.get_docstring(node))
+
+        node = ast.parse('class foo:\n  pass')
+        self.assertIsNone(ast.get_docstring(node.body[0]))
+        node = ast.parse('class foo:\n  x = "not docstring"')
+        self.assertIsNone(ast.get_docstring(node.body[0]))
+        node = ast.parse('class foo:\n  def bar(self): pass')
+        self.assertIsNone(ast.get_docstring(node.body[0]))
+
+        node = ast.parse('def foo():\n  pass')
+        self.assertIsNone(ast.get_docstring(node.body[0]))
+        node = ast.parse('def foo():\n  x = "not docstring"')
+        self.assertIsNone(ast.get_docstring(node.body[0]))
+
+        node = ast.parse('async def foo():\n  pass')
+        self.assertIsNone(ast.get_docstring(node.body[0]))
+        node = ast.parse('async def foo():\n  x = "not docstring"')
+        self.assertIsNone(ast.get_docstring(node.body[0]))
 
     def test_literal_eval(self):
         self.assertEqual(ast.literal_eval('[1, 2, 3]'), [1, 2, 3])
diff --git a/Misc/NEWS.d/next/Library/2018-06-13-15-12-25.bpo-33851.SVbqlz.rst b/Misc/NEWS.d/next/Library/2018-06-13-15-12-25.bpo-33851.SVbqlz.rst
new file mode 100644
index 000000000000..b769ff74a881
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2018-06-13-15-12-25.bpo-33851.SVbqlz.rst
@@ -0,0 +1 @@
+Fix :func:`ast.get_docstring` for a node that lacks a docstring.