def assertTypesMatchPytd(self, ty, pytd_src):
"""Parses pytd_src and compares with ty."""
pytd_tree = parser.parse_string(
textwrap.dedent(pytd_src),
options=parser.PyiOptions(python_version=self.python_version))
pytd_tree = pytd_tree.Visit(
visitors.LookupBuiltins(self.loader.builtins, full_names=False))
pytd_tree = pytd_tree.Visit(visitors.LookupLocalTypes())
pytd_tree = pytd_tree.Visit(visitors.ClassTypeToNamedType())
pytd_tree = pytd_tree.Visit(
visitors.CanonicalOrderingVisitor(sort_signatures=True))
pytd_tree.Visit(visitors.VerifyVisitor())
ty = ty.Visit(visitors.ClassTypeToNamedType())
ty = ty.Visit(visitors.AdjustSelf())
ty = ty.Visit(visitors.CanonicalOrderingVisitor(sort_signatures=True))
ty.Visit(visitors.VerifyVisitor())
ty_src = pytd_utils.Print(ty) + "\n"
pytd_tree_src = pytd_utils.Print(pytd_tree) + "\n"
log.info("========== result ==========")
_LogLines(log.info, ty_src)
log.info("========== expected ==========")
_LogLines(log.info, pytd_tree_src)
log.info("==============================")
# In the diff output, mark expected with "-" and actual with "+".
# (In other words, display a change from "working" to "broken")
self.assertMultiLineEqual(pytd_tree_src, ty_src)
def VisitClass(self, cls):
"""Add superclass methods and constants to this Class."""
if any(base for base in cls.parents
if isinstance(base, pytd.NamedType)):
raise AssertionError("AddInheritedMethods needs a resolved AST")
# Filter out only the types we can reason about.
# TODO(kramm): Do we want handle UnionTypes and GenericTypes at some point?
bases = [
base.cls for base in cls.parents
if isinstance(base, pytd.ClassType)
]
# Don't pull in methods that are named the same as existing methods in
# this class, local methods override parent class methods.
names = {m.name for m in cls.methods} | {c.name for c in cls.constants}
# TODO(kramm): This should do full-blown MRO.
adjust_self = visitors.AdjustSelf(force=True)
adjust_self.class_types.append(visitors.ClassAsType(cls))
new_methods = list(cls.methods)
for base in bases:
for m in base.methods:
if m.name not in names:
new_methods.append(m.Visit(adjust_self))
new_constants = list(cls.constants)
for base in bases:
for c in base.constants:
if c.name not in names:
new_constants.append(c)
return cls.Replace(methods=tuple(new_methods),
constants=tuple(new_constants))
def RemoveMutableParameters(ast):
"""Change all mutable parameters in a pytd AST to a non-mutable form."""
ast = ast.Visit(optimize.AbsorbMutableParameters())
ast = ast.Visit(optimize.CombineContainers())
ast = ast.Visit(optimize.MergeTypeParameters())
ast = ast.Visit(visitors.AdjustSelf(force=True))
return ast
def Optimize(node,
builtins=None,
lossy=False,
use_abcs=False,
max_union=7,
remove_mutable=False,
can_do_lookup=True):
"""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.
builtins: Definitions of all of the external types in node.
lossy: Allow optimizations that change the meaning of the pytd.
use_abcs: Use abstract base classes to represent unions like
e.g. "float or int" as "Real".
max_union: How many types we allow in a union before we simplify
it to just "object".
remove_mutable: Whether to simplify mutable parameters to normal
parameters.
can_do_lookup: True: We're either allowed to try to resolve NamedType
instances in the AST, or the AST is already resolved. False: Skip any
optimizations that would require NamedTypes to be resolved.
Returns:
An optimized node.
"""
node = node.Visit(RemoveDuplicates())
node = node.Visit(SimplifyUnions())
node = node.Visit(CombineReturnsAndExceptions())
node = node.Visit(Factorize())
node = node.Visit(ApplyOptionalArguments())
node = node.Visit(CombineContainers())
node = node.Visit(SimplifyContainers())
if builtins:
superclasses = builtins.Visit(visitors.ExtractSuperClassesByName())
superclasses.update(node.Visit(visitors.ExtractSuperClassesByName()))
if use_abcs:
superclasses.update(abc_hierarchy.GetSuperClasses())
hierarchy = SuperClassHierarchy(superclasses)
node = node.Visit(SimplifyUnionsWithSuperclasses(hierarchy))
if lossy:
node = node.Visit(FindCommonSuperClasses(hierarchy))
if max_union:
node = node.Visit(CollapseLongUnions(max_union))
node = node.Visit(AdjustReturnAndConstantGenericType())
if remove_mutable:
node = node.Visit(AbsorbMutableParameters())
node = node.Visit(CombineContainers())
node = node.Visit(MergeTypeParameters())
node = node.Visit(visitors.AdjustSelf())
node = node.Visit(SimplifyContainers())
if builtins and can_do_lookup:
node = visitors.LookupClasses(node, builtins, ignore_late_types=True)
node = node.Visit(RemoveInheritedMethods())
node = node.Visit(RemoveRedundantSignatures(hierarchy))
return node
def VisitClass(self, cls):
"""Visit a class, and change constants to methods where possible."""
new_constants = []
new_methods = list(cls.methods)
for const in cls.constants:
if self._IsSimpleCall(const.type):
c = self._MaybeLookup(const.type)
signatures = c.methods[0].signatures
self._processed_count[c.name] += 1
new_methods.append(
pytd.Function(const.name, signatures, c.methods[0].kind))
else:
new_constants.append(const) # keep
cls = cls.Replace(constants=tuple(new_constants),
methods=tuple(new_methods))
return cls.Visit(visitors.AdjustSelf(force=True))
def VisitClass(self, cls):
"""Visit a class, and change constants to methods where possible."""
new_constants = []
new_methods = list(cls.methods)
adjust_self = visitors.AdjustSelf(force=True)
adjust_self.class_types.append(visitors.ClassAsType(cls))
for const in cls.constants:
c = self._LookupIfSimpleCall(const.type)
if c:
signatures = c.methods[0].signatures
self._processed_count[c.name] += 1
new_method = pytd.Function(const.name, signatures, c.methods[0].kind)
new_methods.append(new_method.Visit(adjust_self))
else:
new_constants.append(const) # keep
return cls.Replace(constants=tuple(new_constants),
methods=tuple(new_methods))
