def testCollapseLongConstantUnions(self):
src = textwrap.dedent("""
x = ... # type: A or B or C or D
y = ... # type: A or B or C or D or E
""")
expected = textwrap.dedent("""
x = ... # type: A or B or C or D
y = ... # type: ?
""")
ast = self.ParseAndResolve(src)
ast = ast.Visit(optimize.CollapseLongUnions(max_length=4))
ast = ast.Visit(optimize.AdjustReturnAndConstantGenericType())
self.AssertSourceEquals(ast, expected)
def test_collapse_long_constant_unions(self):
src = pytd_src("""
x = ... # type: Union[A, B, C, D]
y = ... # type: Union[A, B, C, D, E]
""")
expected = pytd_src("""
from typing import Any
x = ... # type: Union[A, B, C, D]
y = ... # type: Any
""")
ast = self.ParseAndResolve(src)
ast = ast.Visit(optimize.CollapseLongUnions(max_length=4))
ast = ast.Visit(optimize.AdjustReturnAndConstantGenericType())
self.AssertSourceEquals(ast, expected)
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
""")
new_src = self.ApplyVisitorToString(
src, optimize.CollapseLongUnions(max_length=4))
self.AssertSourceEquals(new_src, expected)
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))
self.AssertSourceEquals(ast, expected)
