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)
