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)
