def Print(ast, print_format=None):
if print_format and print_format not in OUTPUT_FORMATS:
raise ValueError("Invalid format %s" % print_format)
if print_format == "pytd" or print_format is None:
res = ast.Visit(visitors.PrintVisitor())
elif print_format == "pep484stub":
res = ast.Visit(pep484.Print484StubVisitor())
return res
def testPrintImportsNamedType(self):
# Can't get tree by parsing so build explicitly
node = pytd.Constant("x", pytd.NamedType("typing.List"))
tree = pytd.TypeDeclUnit(constants=(node, ),
type_params=(),
functions=(),
classes=(),
aliases=(),
name=None)
expected_src = textwrap.dedent("""
import typing
x = ... # type: typing.List
""").strip()
res = tree.Visit(visitors.PrintVisitor())
self.assertMultiLineEqual(res, expected_src)
def testPrintImports(self):
no_import_src = textwrap.dedent("""
def f(x: Union[int, slice]) -> List[Any]: ...
def g(x: foo.C.C2) -> None: ...
""")
imports = textwrap.dedent("""
import foo.C
from typing import Any, List, Union
""")
expected_src = (imports + no_import_src).strip() # Extra newlines
tree = self.Parse(no_import_src)
res = tree.Visit(visitors.PrintVisitor())
# AssertSourceEquals strips imports
self.AssertSourceEquals(res, no_import_src)
self.assertMultiLineEqual(res, expected_src)
def Print(ast): return ast.Visit(visitors.PrintVisitor())
def testPrintImportsIgnoresExisting(self):
src = "from foo import b"
tree = self.Parse(src)
res = tree.Visit(visitors.PrintVisitor())
self.assertMultiLineEqual(res, src)