Commit 08f127a3 authored by Serhiy Storchaka's avatar Serhiy Storchaka Committed by GitHub

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

parent 9e7c9219
...@@ -206,7 +206,7 @@ def get_docstring(node, clean=True): ...@@ -206,7 +206,7 @@ def get_docstring(node, clean=True):
""" """
if not isinstance(node, (AsyncFunctionDef, FunctionDef, ClassDef, Module)): if not isinstance(node, (AsyncFunctionDef, FunctionDef, ClassDef, Module)):
raise TypeError("%r can't have docstrings" % node.__class__.__name__) 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 return None
node = node.body[0].value node = node.body[0].value
if isinstance(node, Str): if isinstance(node, Str):
...@@ -215,7 +215,7 @@ def get_docstring(node, clean=True): ...@@ -215,7 +215,7 @@ def get_docstring(node, clean=True):
text = node.value text = node.value
else: else:
return None return None
if clean and text: if clean:
import inspect import inspect
text = inspect.cleandoc(text) text = inspect.cleandoc(text)
return text return text
......
...@@ -521,13 +521,44 @@ class ASTHelpers_Test(unittest.TestCase): ...@@ -521,13 +521,44 @@ class ASTHelpers_Test(unittest.TestCase):
) )
def test_get_docstring(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"""') node = ast.parse('def foo():\n """line one\n line two"""')
self.assertEqual(ast.get_docstring(node.body[0]), self.assertEqual(ast.get_docstring(node.body[0]),
'line one\nline two') 'line one\nline two')
node = ast.parse('async def foo():\n """spam\n ham"""') node = ast.parse('async def foo():\n """spam\n ham"""')
self.assertEqual(ast.get_docstring(node.body[0]), 'spam\nham') self.assertEqual(ast.get_docstring(node.body[0]), 'spam\nham')
def test_get_docstring_none(self):
self.assertIsNone(ast.get_docstring(ast.parse(''))) 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): def test_literal_eval(self):
self.assertEqual(ast.literal_eval('[1, 2, 3]'), [1, 2, 3]) self.assertEqual(ast.literal_eval('[1, 2, 3]'), [1, 2, 3])
......
Fix :func:`ast.get_docstring` for a node that lacks a docstring.
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment