pycode: ast.unparse() construct number literals using source code

Developers can write number literals in several ways. For example,
decimal (1234), hexadecimal (0x1234), octal decimal (0o1234) and so on.
But, AST module don't mind how the numbers written in the code. As a
result, ast.unparse() could not reproduce the original form of number
literals.

This allows to construct number literals as possible using original
source code.

Note: This is only available in Python 3.8+.
This commit is contained in:
Takeshi KOMIYA 2020-10-04 01:45:47 +09:00
parent a8abb9995f
commit 0b32e72635
2 changed files with 18 additions and 7 deletions

View File

@ -58,17 +58,19 @@ def parse(code: str, mode: str = 'exec') -> "ast.AST":
return ast.parse(code, mode=mode) return ast.parse(code, mode=mode)
def unparse(node: Optional[ast.AST]) -> Optional[str]: def unparse(node: Optional[ast.AST], code: str = '') -> Optional[str]:
"""Unparse an AST to string.""" """Unparse an AST to string."""
if node is None: if node is None:
return None return None
elif isinstance(node, str): elif isinstance(node, str):
return node return node
return _UnparseVisitor().visit(node) return _UnparseVisitor(code).visit(node)
# a greatly cut-down version of `ast._Unparser` # a greatly cut-down version of `ast._Unparser`
class _UnparseVisitor(ast.NodeVisitor): class _UnparseVisitor(ast.NodeVisitor):
def __init__(self, code: str = '') -> None:
self.code = code
def _visit_op(self, node: ast.AST) -> str: def _visit_op(self, node: ast.AST) -> str:
return OPERATORS[node.__class__] return OPERATORS[node.__class__]
@ -195,6 +197,11 @@ class _UnparseVisitor(ast.NodeVisitor):
def visit_Constant(self, node: ast.Constant) -> str: def visit_Constant(self, node: ast.Constant) -> str:
if node.value is Ellipsis: if node.value is Ellipsis:
return "..." return "..."
elif isinstance(node.value, (int, float, complex)):
if self.code and sys.version_info > (3, 8):
return ast.get_source_segment(self.code, node)
else:
return repr(node.value)
else: else:
return repr(node.value) return repr(node.value)

View File

@ -58,7 +58,7 @@ from sphinx.pycode import ast
]) ])
def test_unparse(source, expected): def test_unparse(source, expected):
module = ast.parse(source) module = ast.parse(source)
assert ast.unparse(module.body[0].value) == expected assert ast.unparse(module.body[0].value, source) == expected
def test_unparse_None(): def test_unparse_None():
@ -66,8 +66,12 @@ def test_unparse_None():
@pytest.mark.skipif(sys.version_info < (3, 8), reason='python 3.8+ is required.') @pytest.mark.skipif(sys.version_info < (3, 8), reason='python 3.8+ is required.')
def test_unparse_py38(): @pytest.mark.parametrize('source,expected', [
source = "lambda x=0, /, y=1, *args, z, **kwargs: x + y + z" ("lambda x=0, /, y=1, *args, z, **kwargs: x + y + z",
expected = "lambda x=0, /, y=1, *args, z, **kwargs: ..." "lambda x=0, /, y=1, *args, z, **kwargs: ..."), # posonlyargs
("0x1234", "0x1234"), # Constant
("1_000_000", "1_000_000"), # Constant
])
def test_unparse_py38(source, expected):
module = ast.parse(source) module = ast.parse(source)
assert ast.unparse(module.body[0].value) == expected assert ast.unparse(module.body[0].value, source) == expected