def testCollapseLongUnions(self):
src = textwrap.dedent("""
def f(x: A or B or C or D) -> X
def g(x: A or B or C or D or E) -> X
def h(x: A or ?) -> X
""")
expected = textwrap.dedent("""
def f(x: A or B or C or D) -> X
def g(x) -> X
def h(x) -> X
""")
ast = self.ParseAndResolve(src)
ast = ast.Visit(optimize.CollapseLongUnions(max_length=4))
ast = ast.Visit(visitors.DropBuiltinPrefix())
self.AssertSourceEquals(ast, expected)
def test_collapse_long_unions(self):
src = pytd_src("""
from typing import Any
def f(x: Union[A, B, C, D]) -> X: ...
def g(x: Union[A, B, C, D, E]) -> X: ...
def h(x: Union[A, Any]) -> X: ...
""")
expected = pytd_src("""
def f(x: Union[A, B, C, D]) -> X: ...
def g(x) -> X: ...
def h(x) -> X: ...
""")
ast = self.ParseAndResolve(src)
ast = ast.Visit(optimize.CollapseLongUnions(max_length=4))
ast = ast.Visit(visitors.DropBuiltinPrefix())
self.AssertSourceEquals(ast, expected)
def testBuiltinSuperClasses(self):
src = textwrap.dedent("""
def f(x: list or object, y: complex or memoryview) -> int or bool
""")
expected = textwrap.dedent("""
def f(x: object, y: object) -> int
""")
hierarchy = self.builtins.Visit(visitors.ExtractSuperClassesByName())
hierarchy.update(self.typing.Visit(visitors.ExtractSuperClassesByName()))
visitor = optimize.FindCommonSuperClasses(
optimize.SuperClassHierarchy(hierarchy))
ast = self.ParseAndResolve(src)
ast = ast.Visit(visitor)
ast = ast.Visit(visitors.DropBuiltinPrefix())
ast = ast.Visit(visitors.CanonicalOrderingVisitor())
self.AssertSourceEquals(ast, expected)
def testRemoveUnknownClasses(self):
src = textwrap.dedent("""
class `~unknown1`():
pass
class `~unknown2`():
pass
class A(object):
def foobar(x: `~unknown1`, y: `~unknown2`) -> `~unknown1` or int
""")
expected = textwrap.dedent("""
class A(object):
def foobar(x, y) -> ? or int
""")
tree = self.Parse(src)
tree = tree.Visit(visitors.RemoveUnknownClasses())
tree = tree.Visit(visitors.DropBuiltinPrefix())
self.AssertSourceEquals(tree, expected)
def test_builtin_superclasses(self):
src = pytd_src("""
def f(x: Union[list, object], y: Union[complex, memoryview]) -> Union[int, bool]: ...
""")
expected = pytd_src("""
def f(x: object, y: object) -> int: ...
""")
hierarchy = self.builtins.Visit(visitors.ExtractSuperClassesByName())
hierarchy.update(
self.typing.Visit(visitors.ExtractSuperClassesByName()))
visitor = optimize.FindCommonSuperClasses(
optimize.SuperClassHierarchy(hierarchy))
ast = self.ParseAndResolve(src)
ast = ast.Visit(visitor)
ast = ast.Visit(visitors.DropBuiltinPrefix())
ast = ast.Visit(visitors.CanonicalOrderingVisitor())
self.AssertSourceEquals(ast, expected)
def test_remove_unknown_classes(self):
src = pytd_src("""
from typing import Union
class `~unknown1`():
pass
class `~unknown2`():
pass
class A:
def foobar(x: `~unknown1`, y: `~unknown2`) -> Union[`~unknown1`, int]: ...
""")
expected = textwrap.dedent("""
from typing import Any, Union
class A:
def foobar(x, y) -> Union[Any, int]: ...
""")
tree = self.Parse(src)
tree = tree.Visit(visitors.RemoveUnknownClasses())
tree = tree.Visit(visitors.DropBuiltinPrefix())
self.AssertSourceEquals(tree, expected)
def testDuplicatesInUnions(self):
src = textwrap.dedent("""
def a(x: int or float or complex) -> bool
def b(x: int or float) -> bool
def c(x: int or int or int) -> bool
def d(x: int or int) -> bool
def e(x: float or int or int or float) -> bool
def f(x: float or int) -> bool
""")
new_src = textwrap.dedent("""
def a(x) -> bool # max_union=2 makes this object
def b(x: int or float) -> bool
def c(x: int) -> bool
def d(x: int) -> bool
def e(x: float or int) -> bool
def f(x: float or int) -> bool
""")
ast = self.ParseAndResolve(src)
optimized = self.Optimize(ast, lossy=False, max_union=2)
optimized = optimized.Visit(visitors.DropBuiltinPrefix())
self.AssertSourceEquals(optimized, new_src)
def test_duplicates_in_unions(self):
src = pytd_src("""
def a(x: Union[int, float, complex]) -> bool: ...
def b(x: Union[int, float]) -> bool: ...
def c(x: Union[int, int, int]) -> bool: ...
def d(x: Union[int, int]) -> bool: ...
def e(x: Union[float, int, int, float]) -> bool: ...
def f(x: Union[float, int]) -> bool: ...
""")
new_src = pytd_src("""
def a(x) -> bool: ... # max_union=2 makes this object
def b(x: Union[int, float]) -> bool: ...
def c(x: int) -> bool: ...
def d(x: int) -> bool: ...
def e(x: Union[float, int]) -> bool: ...
def f(x: Union[float, int]) -> bool: ...
""")
ast = self.ParseAndResolve(src)
optimized = self.Optimize(ast, lossy=False, max_union=2)
optimized = optimized.Visit(visitors.DropBuiltinPrefix())
self.AssertSourceEquals(optimized, new_src)
