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)