def Optimize(node, flags=None):
"""Optimize a PYTD tree.
Tries to shrink a PYTD tree by applying various optimizations.
Arguments:
node: A pytd node to be optimized. It won't be modified - this function will
return a new node.
flags: An instance of OptimizeFlags, to control which optimizations
happen and what parameters to use for the ones that take parameters. Can
be None, in which case defaults will be applied.
Returns:
An optimized node.
"""
node = node.Visit(RemoveDuplicates())
node = node.Visit(CombineReturnsAndExceptions())
node = node.Visit(Factorize())
node = node.Visit(ApplyOptionalArguments())
node = node.Visit(CombineContainers())
if flags and flags.lossy:
hierarchy = node.Visit(visitors.ExtractSuperClassesByName())
node = node.Visit(
FindCommonSuperClasses(hierarchy, flags and flags.use_abcs))
if flags and flags.max_union:
node = node.Visit(CollapseLongParameterUnions(flags.max_union))
node = node.Visit(CollapseLongReturnUnions(flags.max_union))
if flags and flags.remove_mutable:
node = node.Visit(AbsorbMutableParameters())
node = node.Visit(CombineContainers())
node = node.Visit(MergeTypeParameters())
node = node.Visit(visitors.AdjustSelf(force=True))
node = visitors.LookupClasses(node, builtins.GetBuiltins())
node = node.Visit(RemoveInheritedMethods())
return node
def testSuperClassesByName(self):
src = textwrap.dedent("""
class A(nothing):
pass
class B(nothing):
pass
class C(A):
pass
class D(A,B):
pass
class E(C,D,A):
pass
""")
tree = self.Parse(src)
data = tree.Visit(visitors.ExtractSuperClassesByName())
self.assertItemsEqual((), data["A"])
self.assertItemsEqual((), data["B"])
self.assertItemsEqual(("A", ), data["C"])
self.assertItemsEqual(("A", "B"), data["D"])
self.assertItemsEqual(("A", "C", "D"), data["E"])
def testUserSuperClassHierarchy(self):
class_data = textwrap.dedent("""
class AB:
pass
class EFG:
pass
class A(AB, EFG):
pass
class B(AB):
pass
class E(EFG, AB):
pass
class F(EFG):
pass
class G(EFG):
pass
""")
src = textwrap.dedent("""
def f(x: A or B, y: A, z: B) -> E or F or G
def g(x: E or F or G or B) -> E or F
def h(x) -> ?
""") + class_data
expected = textwrap.dedent("""
def f(x: AB, y: A, z: B) -> EFG
def g(x) -> EFG
def h(x) -> ?
""") + class_data
hierarchy = self.Parse(src).Visit(visitors.ExtractSuperClassesByName())
visitor = optimize.FindCommonSuperClasses(hierarchy, use_abcs=False)
new_src = self.ApplyVisitorToString(src, visitor)
self.AssertSourceEquals(new_src, expected)
def GetBuiltinsHierarchy():
builtins = GetBuiltins()
return builtins.Visit(visitors.ExtractSuperClassesByName())