bpo-33851: Fix ast.get_docstring() for a node that lacks a docstring. (GH-7682)
(cherry picked from commit 08f127a3cad8ce4eb281d30d9488c91b0fd7cfed)
Co-authored-by: Serhiy Storchaka <storchaka@gmail.com>
diff --git a/Lib/ast.py b/Lib/ast.py
index 134d9d2..bfe346b 100644
--- a/Lib/ast.py
+++ b/Lib/ast.py
@@ -206,7 +206,7 @@
"""
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 @@
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 ab32d9d..7db40e7 100644
--- a/Lib/test/test_ast.py
+++ b/Lib/test/test_ast.py
@@ -521,13 +521,44 @@
)
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])