diff options
Diffstat (limited to 'Lib')
| -rw-r--r-- | Lib/ast.py | 9 | ||||
| -rw-r--r-- | Lib/test/test_unparse.py | 14 |
2 files changed, 20 insertions, 3 deletions
@@ -741,6 +741,15 @@ class _Unparser(NodeVisitor): 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) diff --git a/Lib/test/test_unparse.py b/Lib/test/test_unparse.py index 3d87cfb..2329264 100644 --- a/Lib/test/test_unparse.py +++ b/Lib/test/test_unparse.py @@ -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): @@ -330,6 +330,14 @@ class UnparseTestCase(ASTTestCase): 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""" |