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 object) -> 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 testBuiltinSuperClasses(self):
src = textwrap.dedent("""
def f(x: list or object, y: int or float) -> int or bool
""")
expected = textwrap.dedent("""
def f(x, y) -> int
""")
b = builtins.GetBuiltinsPyTD()
hierarchy = b.Visit(visitors.ExtractSuperClassesByName())
visitor = optimize.FindCommonSuperClasses(
optimize.SuperClassHierarchy(hierarchy))
ast = self.ParseAndResolve(src)
ast = ast.Visit(visitor)
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, y) -> int
""")
b = builtins.GetBuiltinsPyTD(self.PYTHON_VERSION)
hierarchy = b.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 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 = optimize.Optimize(ast, lossy=False, max_union=2)
optimized = optimized.Visit(visitors.DropBuiltinPrefix())
self.AssertSourceEquals(optimized, new_src)
