def testExpandCompatibleBuiltins(self):
src = textwrap.dedent("""
from typing import Tuple, Union
def f1(a: float) -> None: ...
def f2() -> float: ...
def f3(a: bool) -> None: ...
def f4() -> bool: ...
def f5(a: Union[bool, int]) -> None: ...
def f6(a: Tuple[bool, int]) -> None: ...
""")
expected = textwrap.dedent("""
from typing import Tuple, Union
def f1(a: Union[float, int]) -> None: ...
def f2() -> float: ...
def f3(a: Union[bool, None]) -> None: ...
def f4() -> bool: ...
def f5(a: Union[bool, None, int]) -> None: ...
def f6(a: Tuple[Union[bool, None], int]) -> None: ...
""")
src_tree, expected_tree = (self.Parse(s).Visit(
visitors.LookupBuiltins(self.loader.builtins))
for s in (src, expected))
new_tree = src_tree.Visit(
visitors.ExpandCompatibleBuiltins(self.loader.builtins))
self.AssertSourceEquals(new_tree, expected_tree)
def _resolve_builtins(self, pyval, ast=None):
builtins_lookup = visitors.LookupBuiltins(self.builtins, full_names=False)
if ast:
builtins_lookup.EnterTypeDeclUnit(ast)
pyval = pyval.Visit(builtins_lookup)
pyval = pyval.Visit(visitors.ExpandCompatibleBuiltins(self.builtins))
return pyval
def resolve_builtin_types(self, mod_ast, *, lookup_ast=None):
bltn_lookup = visitors.LookupBuiltins(self.builtins_ast,
full_names=False)
mod_ast = self._lookup(bltn_lookup, mod_ast, lookup_ast)
mod_ast = mod_ast.Visit(
visitors.ExpandCompatibleBuiltins(self.builtins_ast))
return mod_ast
def _postprocess_pyi(self, ast):
"""Apply all the PYI transformations we need."""
ast = ast.Visit(visitors.LookupBuiltins(self.builtins, full_names=False))
ast = ast.Visit(visitors.ExpandCompatibleBuiltins(self.builtins))
dependencies = self._collect_ast_dependencies(ast)
if dependencies:
self._load_ast_dependencies(dependencies, ast)
ast = self._resolve_external_types(ast)
ast = ast.Visit(visitors.LookupLocalTypes())
return ast
def _postprocess_pyi(self, ast):
"""Apply all the PYI transformations we need."""
package_name = module_utils.get_package_name(ast.name, ast.is_package)
if package_name:
ast = ast.Visit(visitors.QualifyRelativeNames(package_name))
ast = ast.Visit(
visitors.LookupBuiltins(self.builtins, full_names=False))
ast = ast.Visit(visitors.ExpandCompatibleBuiltins(self.builtins))
dependencies = self._collect_ast_dependencies(ast)
if dependencies:
self._load_ast_dependencies(dependencies, ast)
ast = self._resolve_external_types(ast)
ast = ast.Visit(visitors.LookupLocalTypes())
return ast
def PrepareForExport(module_name, python_version, ast, loader):
"""Prepare an ast as if it was parsed and loaded.
External dependencies will not be resolved, as the ast generated by this
method is supposed to be exported.
Args:
module_name: The module_name as a string for the returned ast.
python_version: A tuple of (major, minor) python version as string
(see config.python_version).
ast: pytd.TypeDeclUnit, is only used if src is None.
loader: A load_pytd.Loader instance.
Returns:
A pytd.TypeDeclUnit representing the supplied AST as it would look after
being written to a file and parsed.
"""
# This is a workaround for functionality which crept into places it doesn't
# belong. Ideally this would call some transformation Visitors on ast to
# transform it into the same ast we get after parsing and loading (compare
# load_pytd.Loader.load_file). Unfortunately parsing has some special cases,
# e.g. '__init__' return type and '__new__' being a 'staticmethod', which
# need to be moved to visitors before we can do this. Printing an ast also
# applies transformations,
# e.g. visitors.PrintVisitor._FormatContainerContents, which need to move to
# their own visitors so they can be applied without printing.
src = pytd_utils.Print(ast)
ast = parser.parse_string(src=src,
name=module_name,
python_version=python_version)
ast = ast.Visit(visitors.LookupBuiltins(loader.builtins, full_names=False))
ast = ast.Visit(
visitors.ExpandCompatibleBuiltins(loader.builtins, python_version))
ast = ast.Visit(visitors.LookupLocalTypes())
ast = ast.Visit(visitors.AdjustTypeParameters())
ast = ast.Visit(visitors.NamedTypeToClassType())
ast = ast.Visit(visitors.FillInLocalPointers({"": ast, module_name: ast}))
ast = ast.Visit(visitors.CanonicalOrderingVisitor())
ast = ast.Visit(
visitors.ClassTypeToLateType(
ignore=[module_name + ".", "__builtin__.", "typing."]))
return ast
def _resolve_builtins(self, ast):
ast = ast.Visit(
visitors.LookupBuiltins(self.builtins, full_names=False))
ast = ast.Visit(visitors.ExpandCompatibleBuiltins(self.builtins))
return ast
