Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions Lib/ast.py
Original file line number Diff line number Diff line change
Expand Up @@ -741,6 +741,15 @@ def _write_docstring_and_traverse_body(self, node):
def visit_Module(self, node):
self._write_docstring_and_traverse_body(node)

def visit_FunctionType(self, node):
with self.delimit("(", ")"):
self.interleave(
lambda: self.write(", "), self.traverse, node.argtypes
)

self.write(" -> ")
self.traverse(node.returns)

def visit_Expr(self, node):
self.fill()
self.set_precedence(_Precedence.YIELD, node.value)
Expand Down
14 changes: 11 additions & 3 deletions Lib/test/test_unparse.py
Original file line number Diff line number Diff line change
Expand Up @@ -122,10 +122,10 @@ class ASTTestCase(unittest.TestCase):
def assertASTEqual(self, ast1, ast2):
self.assertEqual(ast.dump(ast1), ast.dump(ast2))

def check_ast_roundtrip(self, code1):
ast1 = ast.parse(code1)
def check_ast_roundtrip(self, code1, **kwargs):
ast1 = ast.parse(code1, **kwargs)
code2 = ast.unparse(ast1)
ast2 = ast.parse(code2)
ast2 = ast.parse(code2, **kwargs)
self.assertASTEqual(ast1, ast2)

def check_invalid(self, node, raises=ValueError):
Expand Down Expand Up @@ -330,6 +330,14 @@ def test_constant_tuples(self):
ast.Constant(value=(1, 2, 3), kind=None), "(1, 2, 3)"
)

def test_function_type(self):
for function_type in (
"() -> int",
"(int, int) -> int",
"(Callable[complex], More[Complex(call.to_typevar())]) -> None"
):
self.check_ast_roundtrip(function_type, mode="func_type")


class CosmeticTestCase(ASTTestCase):
"""Test if there are cosmetic issues caused by unnecesary additions"""
Expand Down