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 load(self, options):
"""Read builtins.pytd and typing.pytd, and return the parsed modules."""
t = self._parse_predefined("typing", options)
b = self._parse_predefined("builtins", options)
b = b.Visit(
visitors.LookupExternalTypes({"typing": t}, self_name="builtins"))
t = t.Visit(visitors.LookupBuiltins(b))
b = b.Visit(visitors.NamedTypeToClassType())
t = t.Visit(visitors.NamedTypeToClassType())
b = b.Visit(visitors.AdjustTypeParameters())
t = t.Visit(visitors.AdjustTypeParameters())
b = b.Visit(visitors.CanonicalOrderingVisitor())
t = t.Visit(visitors.CanonicalOrderingVisitor())
b.Visit(
visitors.FillInLocalPointers({
"": b,
"typing": t,
"builtins": b
}))
t.Visit(
visitors.FillInLocalPointers({
"": t,
"typing": t,
"builtins": b
}))
b.Visit(visitors.VerifyLookup())
t.Visit(visitors.VerifyLookup())
b.Visit(visitors.VerifyContainers())
t.Visit(visitors.VerifyContainers())
return b, t
def testCanonicalOrderingVisitor(self):
src1 = textwrap.dedent("""
from typing import TypeVar
def f() -> ?:
raise MemoryError()
raise IOError()
def f(x: list[a]) -> ?
def f(x: list[b or c]) -> ?
def f(x: list[tuple[d]]) -> ?
A = TypeVar("A")
C = TypeVar("C")
B = TypeVar("B")
D = TypeVar("D")
def f(d: A, c: B, b: C, a: D) -> ?
""")
src2 = textwrap.dedent("""
def f() -> ?:
raise IOError()
raise MemoryError()
def f(x: list[tuple[d]]) -> ?
def f(x: list[a]) -> ?
def f(x: list[b or c]) -> ?
A = TypeVar("A")
C = TypeVar("C")
B = TypeVar("B")
D = TypeVar("D")
def f(d: A, c: B, b: C, a: D) -> ?
""")
tree1 = self.Parse(src1)
tree1 = tree1.Visit(visitors.CanonicalOrderingVisitor(sort_signatures=True))
tree2 = self.Parse(src2)
tree2 = tree2.Visit(visitors.CanonicalOrderingVisitor(sort_signatures=True))
self.AssertSourceEquals(tree1, tree2)
self.assertEqual(tree1.Lookup("f").signatures[0].template,
tree2.Lookup("f").signatures[0].template)
def canonical_pyi(pyi, multiline_args=False, options=None): """Rewrite a pyi in canonical form.""" ast = parse_string(pyi, options=options) ast = ast.Visit(visitors.ClassTypeToNamedType()) ast = ast.Visit(visitors.CanonicalOrderingVisitor(sort_signatures=True)) ast.Visit(visitors.VerifyVisitor()) return pytd_utils.Print(ast, multiline_args)
def StoreAst(ast, filename=None, open_function=open):
"""Loads and stores an ast to disk.
Args:
ast: The pytd.TypeDeclUnit to save to disk.
filename: The filename for the pickled output. If this is None, this
function instead returns the pickled string.
open_function: A custom file opening function.
Returns:
The pickled string, if no filename was given. (None otherwise.)
"""
if ast.name.endswith(".__init__"):
ast = ast.Visit(
visitors.RenameModuleVisitor(ast.name,
ast.name.rsplit(".__init__", 1)[0]))
# Collect dependencies
deps = visitors.CollectDependencies()
ast.Visit(deps)
dependencies = deps.dependencies
late_dependencies = deps.late_dependencies
# Clean external references
ast.Visit(visitors.ClearClassPointers())
indexer = FindClassTypesVisitor()
ast.Visit(indexer)
ast = ast.Visit(visitors.CanonicalOrderingVisitor())
return pytd_utils.SavePickle(SerializableAst(
ast, sorted(dependencies.items()), sorted(late_dependencies.items()),
sorted(indexer.class_type_nodes)),
filename,
open_function=open_function)
def GetBuiltinsAndTyping(python_version): # Deprecated. Use load_pytd instead.
"""Get __builtin__.pytd and typing.pytd."""
assert python_version
global _cached_builtins_pytd
if _cached_builtins_pytd.cache:
assert _cached_builtins_pytd.version == python_version
else:
t = parser.parse_string(_FindBuiltinFile("typing", python_version),
name="typing",
python_version=python_version)
b = parser.parse_string(_FindBuiltinFile("__builtin__",
python_version),
name="__builtin__",
python_version=python_version)
b = b.Visit(
visitors.LookupExternalTypes({"typing": t},
full_names=True,
self_name="__builtin__"))
t = t.Visit(visitors.LookupBuiltins(b))
b = b.Visit(visitors.NamedTypeToClassType())
t = t.Visit(visitors.NamedTypeToClassType())
b = b.Visit(visitors.AdjustTypeParameters())
t = t.Visit(visitors.AdjustTypeParameters())
b = b.Visit(visitors.CanonicalOrderingVisitor())
t = t.Visit(visitors.CanonicalOrderingVisitor())
b.Visit(
visitors.FillInLocalPointers({
"": b,
"typing": t,
"__builtin__": b
}))
t.Visit(
visitors.FillInLocalPointers({
"": t,
"typing": t,
"__builtin__": b
}))
b.Visit(visitors.VerifyLookup())
t.Visit(visitors.VerifyLookup())
b.Visit(visitors.VerifyContainers())
t.Visit(visitors.VerifyContainers())
_cached_builtins_pytd = Cache(python_version, (b, t))
return _cached_builtins_pytd.cache
def GetBuiltinsAndTyping(python_version): # Deprecated. Use load_pytd instead.
"""Get builtins.pytd and typing.pytd."""
assert python_version
if python_version not in _cached_builtins_pytd:
t = parser.parse_string(_FindBuiltinFile("typing", python_version),
name="typing",
python_version=python_version)
b = parser.parse_string(_FindBuiltinFile("builtins", python_version),
name="builtins",
python_version=python_version)
b = b.Visit(
visitors.LookupExternalTypes({"typing": t}, self_name="builtins"))
t = t.Visit(visitors.LookupBuiltins(b))
b = b.Visit(visitors.NamedTypeToClassType())
t = t.Visit(visitors.NamedTypeToClassType())
b = b.Visit(visitors.AdjustTypeParameters())
t = t.Visit(visitors.AdjustTypeParameters())
b = b.Visit(visitors.CanonicalOrderingVisitor())
t = t.Visit(visitors.CanonicalOrderingVisitor())
b.Visit(
visitors.FillInLocalPointers({
"": b,
"typing": t,
"builtins": b
}))
t.Visit(
visitors.FillInLocalPointers({
"": t,
"typing": t,
"builtins": b
}))
b.Visit(visitors.VerifyLookup())
t.Visit(visitors.VerifyLookup())
b.Visit(visitors.VerifyContainers())
t.Visit(visitors.VerifyContainers())
_cached_builtins_pytd[python_version] = (b, t)
return _cached_builtins_pytd[python_version]
def test_builtin_superclasses(self):
src = pytd_src("""
def f(x: Union[list, object], y: Union[complex, memoryview]) -> Union[int, bool]: ...
""")
expected = pytd_src("""
def f(x: builtins.object, y: builtins.object) -> builtins.int: ...
""")
hierarchy = self.builtins.Visit(visitors.ExtractSuperClassesByName())
hierarchy.update(
self.typing.Visit(visitors.ExtractSuperClassesByName()))
visitor = optimize.FindCommonSuperClasses(
optimize.SuperClassHierarchy(hierarchy))
ast = self.ParseAndResolve(src)
ast = ast.Visit(visitor)
ast = ast.Visit(visitors.CanonicalOrderingVisitor())
self.AssertSourceEquals(ast, expected)
def testBuiltinSuperClasses(self):
src = textwrap.dedent("""
def f(x: list or object, y: complex or memoryview) -> int or bool
""")
expected = textwrap.dedent("""
def f(x: object, y: object) -> int
""")
hierarchy = self.builtins.Visit(visitors.ExtractSuperClassesByName())
hierarchy.update(self.typing.Visit(visitors.ExtractSuperClassesByName()))
visitor = optimize.FindCommonSuperClasses(
optimize.SuperClassHierarchy(hierarchy))
ast = self.ParseAndResolve(src)
ast = ast.Visit(visitor)
ast = ast.Visit(visitors.DropBuiltinPrefix())
ast = ast.Visit(visitors.CanonicalOrderingVisitor())
self.AssertSourceEquals(ast, expected)
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 _get_ast(self, temp_dir, module_name, src=None):
src = src or ("""
import module2
from module2 import f
from typing import List
constant = True
x = List[int]
b = List[int]
class SomeClass(object):
def __init__(self, a: module2.ObjectMod2):
pass
def ModuleFunction():
pass
""")
pyi_filename = temp_dir.create_file("module1.pyi", src)
temp_dir.create_file(
"module2.pyi", """
import queue
def f() -> queue.Queue: ...
class ObjectMod2(object):
def __init__(self):
pass
""")
loader = load_pytd.Loader(base_module=None,
python_version=self.python_version,
pythonpath=[temp_dir.path])
ast = loader.load_file(module_name, pyi_filename)
# serialize_ast.StoreAst sorts the ast for determinism, so we should do the
# same to the original ast to do pre- and post-pickling comparisons.
loader._modules[module_name].ast = ast = ast.Visit(
visitors.CanonicalOrderingVisitor())
return ast, loader
def canonical_pyi(pyi, python_version, multiline_args=False):
ast = parser.parse_string(pyi, python_version=python_version)
ast = ast.Visit(visitors.ClassTypeToNamedType())
ast = ast.Visit(visitors.CanonicalOrderingVisitor(sort_signatures=True))
ast.Visit(visitors.VerifyVisitor())
return Print(ast, multiline_args)
def CanonicalOrdering(n, sort_signatures=False):
"""Convert a PYTD node to a canonical (sorted) ordering."""
return n.Visit(
visitors.CanonicalOrderingVisitor(sort_signatures=sort_signatures))
def assertPytd(self, val, expected):
pytd_tree = val.to_type()
pytd_tree = pytd_tree.Visit(visitors.CanonicalOrderingVisitor())
actual = pytd_utils.Print(pytd_tree)
self.assertEqual(actual, expected)
