def testFactorize(self): src = textwrap.dedent(""" def foo(a: int) -> file def foo(a: int, x: complex) -> file def foo(a: int, x: str) -> file def foo(a: float, x: complex) -> file def foo(a: float, x: str) -> file def foo(a: int, x: file, ...) -> file """) new_src = textwrap.dedent(""" def foo(a: int) -> file def foo(a: float, x: complex or str) -> file def foo(a: int, x: file, ...) -> file """) self.AssertSourceEquals( self.ApplyVisitorToString(src, optimize.Factorize()), new_src)
def test_factorize(self): src = pytd_src(""" def foo(a: int) -> file: ... def foo(a: int, x: complex) -> file: ... def foo(a: int, x: str) -> file: ... def foo(a: float, x: complex) -> file: ... def foo(a: float, x: str) -> file: ... def foo(a: int, x: file, *args) -> file: ... """) new_src = pytd_src(""" def foo(a: int) -> file: ... def foo(a: float, x: Union[complex, str]) -> file: ... def foo(a: int, x: file, *args) -> file: ... """) self.AssertSourceEquals( self.ApplyVisitorToString(src, optimize.Factorize()), new_src)
def testFactorizeMutable(self): src = textwrap.dedent(""" def foo(a: list[bool], b: X) -> file: a := list[int] def foo(a: list[bool], b: Y) -> file: a := list[int] # not groupable: def bar(a: int, b: list[int]) -> file: b := list[complex] def bar(a: int, b: list[float]) -> file: b := list[str] """) new_src = textwrap.dedent(""" def foo(a: list[bool], b: X or Y) -> file: a := list[int] def bar(a: int, b: list[int]) -> file: b := list[complex] def bar(a: int, b: list[float]) -> file: b := list[str] """) self.AssertSourceEquals( self.ApplyVisitorToString(src, optimize.Factorize()), new_src)
def test_factorize_mutable(self): src = pytd_src(""" def foo(a: list[bool], b: X) -> file: a = list[int] def foo(a: list[bool], b: Y) -> file: a = list[int] # not groupable: def bar(a: int, b: list[int]) -> file: b = list[complex] def bar(a: int, b: list[float]) -> file: b = list[str] """) new_src = pytd_src(""" def foo(a: list[bool], b: Union[X, Y]) -> file: a = list[int] def bar(a: int, b: list[int]) -> file: b = list[complex] def bar(a: int, b: list[float]) -> file: b = list[str] """) self.AssertSourceEquals( self.ApplyVisitorToString(src, optimize.Factorize()), new_src)