def assertBuiltinsPickleEqual(self, f1, f2):
with open(f1, "rb") as pickle1, open(f2, "rb") as pickle2:
if pickle1.read() == pickle2.read():
return
out1 = pytd_utils.LoadPickle(f1, compress=True)
out2 = pytd_utils.LoadPickle(f2, compress=True)
raise AssertionError("\n".join(pytd_utils.DiffNamedPickles(out1, out2)))
def assertBuiltinsPickleEqual(self, f1, f2):
with open(f1, "rb") as pickle1, open(f2, "rb") as pickle2:
if pickle1.read() == pickle2.read():
return
out1 = pytd_utils.LoadPickle(f1, compress=True)
out2 = pytd_utils.LoadPickle(f2, compress=True)
diff = []
for (name1, ast1), (name2, ast2) in zip(out1, out2):
if name1 != name2:
raise AssertionError("different ordering of pyi files")
elif ast1 != ast2:
diff.append(name1)
raise AssertionError("pyi files differ: " + ", ".join(diff))
def testLoadPickleFromFile(self):
d1 = {1, 2j, "3"}
with utils.Tempdir() as d:
filename = d.create_file("foo.pickle")
pytd_utils.SavePickle(d1, filename)
d2 = pytd_utils.LoadPickle(filename)
self.assertEqual(d1, d2)
def test_load_pickle_from_compressed_file(self):
d1 = {1, 2j, "3"}
with file_utils.Tempdir() as d:
filename = d.create_file("foo.pickle.gz")
pytd_utils.SavePickle(d1, filename, compress=True)
d2 = pytd_utils.LoadPickle(filename, compress=True)
self.assertEqual(d1, d2)
def load_from_pickle(cls, filename, base_module, **kwargs):
items = pytd_utils.LoadPickle(filename, compress=True)
modules = {
name: Module(name, filename=None, ast=None, pickle=pickle, dirty=False)
for name, pickle in items
}
return cls(base_module=base_module, modules=modules, **kwargs)
def load_file(self, module_name, filename, ast=None):
"""Load (or retrieve from cache) a module and resolve its dependencies."""
if not is_pickle(filename):
return super(PickledPyiLoader, self).load_file(module_name, filename, ast)
existing = self._get_existing_ast(module_name)
if existing:
# TODO(kramm): When does this happen?
return existing
loaded_ast = pytd_utils.LoadPickle(filename)
# At this point ast.name and module_name could be different.
# They are later synced in ProcessAst.
dependencies = {d: names for d, names in loaded_ast.dependencies
if d != loaded_ast.ast.name}
loaded_ast = serialize_ast.EnsureAstName(loaded_ast, module_name, fix=True)
self._modules[module_name] = Module(module_name, filename, loaded_ast.ast)
self._load_ast_dependencies(dependencies, ast, module_name)
try:
ast = serialize_ast.ProcessAst(loaded_ast, self._get_module_map())
except serialize_ast.UnrestorableDependencyError as e:
del self._modules[module_name]
raise BadDependencyError(utils.message(e), module_name)
# Mark all the module's late dependencies as explicitly imported.
for d, _ in loaded_ast.late_dependencies:
if d != loaded_ast.ast.name:
self.add_module_prefixes(d)
self._modules[module_name].ast = ast
self._modules[module_name].pickle = None
self._modules[module_name].dirty = False
return ast
def load_file(self, module_name, filename, ast=None):
"""Load (or retrieve from cache) a module and resolve its dependencies."""
if not os.path.splitext(filename)[1].startswith(".pickled"):
return super(PickledPyiLoader,
self).load_file(module_name, filename, ast)
existing = self._get_existing_ast(module_name)
if existing:
# TODO(kramm): When does this happen?
return existing
loaded_ast = pytd_utils.LoadPickle(filename)
# At this point ast.name and module_name could be different.
# They are later synced in ProcessAst.
dependencies = [
d for d in loaded_ast.dependencies if d != loaded_ast.ast.name
]
loaded_ast = serialize_ast.EnsureAstName(loaded_ast,
module_name,
fix=True)
self._modules[module_name] = Module(module_name, filename,
loaded_ast.ast)
self._load_ast_dependencies(dependencies, ast, module_name)
try:
ast = serialize_ast.ProcessAst(loaded_ast, self._get_module_map())
except serialize_ast.UnrestorableDependencyError as e:
del self._modules[module_name]
raise BadDependencyError(e.message, module_name)
self._modules[module_name].ast = ast
self._modules[module_name].pickle = None
self._modules[module_name].dirty = False
return ast
def LoadPrecompiled(filename, python_version):
"""Load precompiled builtins from the specified file."""
# TODO(steenbuck): This should check that the python_version in the current
# process is the same as the one used to generate the cached file.
global _cached_builtins_pytd
assert _cached_builtins_pytd.cache is None
_cached_builtins_pytd = Cache(python_version,
pytd_utils.LoadPickle(filename))
def testUnrestorableDependencyErrorWithModuleIndex(self):
with utils.Tempdir() as d:
module_name = "module1"
pickled_ast_filename = os.path.join(d.path, "module1.pyi.pickled")
module_map = self._StoreAst(d, module_name, pickled_ast_filename)
module_map = {} # Remove module2
with self.assertRaises(serialize_ast.UnrestorableDependencyError):
serialize_ast.ProcessAst(
pytd_utils.LoadPickle(pickled_ast_filename), module_map)
def test_function_type(self):
with file_utils.Tempdir() as d:
foo = d.create_file("foo.pickle")
module_map = self._store_ast(d, "foo", foo, ast=self._get_ast(d, "foo"))
p = pytd_utils.LoadPickle(foo)
self.assertTrue(p.function_type_nodes)
ast = serialize_ast.ProcessAst(p, module_map)
f, = [a for a in ast.aliases if a.name == "foo.f"]
signature, = f.type.function.signatures
self.assertIsNotNone(signature.return_type.cls)
def testUnrestorableDependencyErrorWithoutModuleIndex(self):
with utils.Tempdir() as d:
module_name = "module1"
pickled_ast_filename = os.path.join(d.path, "module1.pyi.pickled")
module_map = self._StoreAst(d, module_name, pickled_ast_filename)
module_map = {} # Remove module2
loaded_ast = pytd_utils.LoadPickle(pickled_ast_filename)
loaded_ast.modified_class_types = None # Remove the index
with self.assertRaises(serialize_ast.UnrestorableDependencyError):
serialize_ast.ProcessAst(loaded_ast, module_map)
def load_from_pickle(cls, filename, base_module, **kwargs):
"""Load a pytd module from a pickle file."""
items = pytd_utils.LoadPickle(
filename, compress=True,
open_function=kwargs.get("open_function", open))
modules = {
name: Module(name, filename=None, ast=None, pickle=pickle,
has_unresolved_pointers=False)
for name, pickle in items
}
return cls(base_module=base_module, modules=modules, **kwargs)
def load_from_pickle(cls, filename, options):
"""Load a pytd module from a pickle file."""
items = pytd_utils.LoadPickle(filename,
compress=True,
open_function=options.open_function)
modules = {
name: Module(name,
filename=None,
ast=None,
pickle=pickle,
has_unresolved_pointers=False)
for name, pickle in items
}
return cls(options, modules=modules)
def testStoreRemovesInit(self):
with utils.Tempdir() as d:
original_module_name = "module1.__init__"
pickled_ast_filename = os.path.join(d.path, "module1.pyi.pickled")
module_map = self._StoreAst(d, original_module_name, pickled_ast_filename)
serializable_ast = pytd_utils.LoadPickle(pickled_ast_filename)
expected_name = "module1"
# Check that the module had the expected name before.
self.assertTrue(original_module_name in module_map)
# Check that module1 wasn't created before storing.
self.assertTrue(expected_name not in module_map)
# Check that the saved ast had its name changed.
self.assertEqual(serializable_ast.ast.name, expected_name)
def testLoadTopLevel(self):
"""Tests that a pickled file can be read."""
with utils.Tempdir() as d:
module_name = "module1"
pickled_ast_filename = os.path.join(d.path, "module1.pyi.pickled")
module_map = self._StoreAst(d, module_name, pickled_ast_filename)
original_ast = module_map[module_name]
del module_map[module_name]
loaded_ast = serialize_ast.ProcessAst(
pytd_utils.LoadPickle(pickled_ast_filename), module_map)
self.assertTrue(loaded_ast)
self.assertTrue(loaded_ast is not original_ast)
self.assertEqual(loaded_ast.name, module_name)
self.assertTrue(original_ast.ASTeq(loaded_ast))
loaded_ast.Visit(visitors.VerifyLookup())
def testLoadWithDifferentModuleName(self):
with utils.Tempdir() as d:
original_module_name = "module1"
pickled_ast_filename = os.path.join(d.path, "module1.pyi.pickled")
module_map = self._StoreAst(d, original_module_name, pickled_ast_filename)
original_ast = module_map[original_module_name]
del module_map[original_module_name]
new_module_name = "wurstbrot.module2"
serializable_ast = pytd_utils.LoadPickle(pickled_ast_filename)
serializable_ast = serialize_ast.EnsureAstName(
serializable_ast, new_module_name)
loaded_ast = serialize_ast.ProcessAst(serializable_ast, module_map)
self.assertTrue(loaded_ast)
self.assertTrue(loaded_ast is not original_ast)
self.assertEqual(loaded_ast.name, new_module_name)
loaded_ast.Visit(visitors.VerifyLookup())
self.assertFalse(original_ast.ASTeq(loaded_ast))
ast_new_module, _ = self._GetAst(temp_dir=d, module_name=new_module_name)
self.assertTrue(ast_new_module.ASTeq(loaded_ast))
def testNodeIndexVisitorUsage(self):
"""Confirms that the node index is used.
This removes the first node from the class_type_nodes list and checks that
that node is not updated by ProcessAst.
"""
with utils.Tempdir() as d:
module_name = "module1"
pickled_ast_filename = os.path.join(d.path, "module1.pyi.pickled")
module_map = self._StoreAst(d, module_name, pickled_ast_filename)
del module_map[module_name]
serialized_ast = pytd_utils.LoadPickle(pickled_ast_filename)
# The sorted makes the testcase more deterministic.
serialized_ast = serialized_ast.Replace(
class_type_nodes=sorted(serialized_ast.class_type_nodes)[1:])
loaded_ast = serialize_ast.ProcessAst(serialized_ast, module_map)
with self.assertRaisesRegexp(
ValueError, "Unresolved class: '__builtin__.NoneType'"):
loaded_ast.Visit(visitors.VerifyLookup())
def testLoadWithSameModuleName(self):
"""Explicitly set the module name and reload with the same name.
The difference to testLoadTopLevel is that the module name does not match
the filelocation.
"""
with utils.Tempdir() as d:
module_name = "foo.bar.module1"
pickled_ast_filename = os.path.join(d.path, "module1.pyi.pickled")
module_map = self._StoreAst(d, module_name, pickled_ast_filename)
original_ast = module_map[module_name]
del module_map[module_name]
loaded_ast = serialize_ast.ProcessAst(
pytd_utils.LoadPickle(pickled_ast_filename), module_map)
self.assertTrue(loaded_ast)
self.assertTrue(loaded_ast is not original_ast)
self.assertEqual(loaded_ast.name, "foo.bar.module1")
self.assertTrue(original_ast.ASTeq(loaded_ast))
loaded_ast.Visit(visitors.VerifyLookup())
def load_file(self, module_name, filename, mod_ast=None):
"""Load (or retrieve from cache) a module and resolve its dependencies."""
if not pytd_utils.IsPickle(filename):
return super().load_file(module_name, filename, mod_ast)
existing = self._modules.get_existing_ast(module_name)
if existing:
return existing
loaded_ast = pytd_utils.LoadPickle(
filename, open_function=self.options.open_function)
# At this point ast.name and module_name could be different.
# They are later synced in ProcessAst.
dependencies = {
d: names
for d, names in loaded_ast.dependencies if d != loaded_ast.ast.name
}
loaded_ast = serialize_ast.EnsureAstName(loaded_ast,
module_name,
fix=True)
self._modules[module_name] = Module(module_name, filename,
loaded_ast.ast)
self._load_ast_dependencies(dependencies,
lookup_ast=mod_ast,
lookup_ast_name=module_name)
try:
ast = serialize_ast.ProcessAst(loaded_ast,
self._modules.get_module_map())
except serialize_ast.UnrestorableDependencyError as e:
del self._modules[module_name]
raise BadDependencyError(utils.message(e), module_name) from e
# Mark all the module's late dependencies as explicitly imported.
for d, _ in loaded_ast.late_dependencies:
if d != loaded_ast.ast.name:
self.add_module_prefixes(d)
self._modules[module_name].ast = ast
self._modules[module_name].pickle = None
self._modules[module_name].has_unresolved_pointers = False
return ast
def testUnrestorableChild(self):
# Assume .cls in a ClassType X in module1 was referencing something for
# which, Visitors.LookupExternalTypes returned AnythingType.
# Now if something in module1 is referencing X external types need to be
# resolved before local types, so that we can resolve local types to the
# correct ClassType, as the ClassType instance changes, if .cls can not be
# filled and instead AnythingType is used.
class RenameVisitor(visitors.Visitor):
def __init__(self, *args, **kwargs):
super(RenameVisitor, self).__init__(*args, **kwargs)
self._init = False
def EnterFunction(self, func):
if func.name == "__init__":
self._init = True
return None
return False
def LeaveFunction(self, func):
self._init = False
def VisitClassType(self, cls_type):
if self._init:
cls_type = cls_type.Replace(
name="other_module.unknown_Reference")
# Needs to be copied manually as it is not part of the NamedTuple.
cls_type.cls = None
return cls_type
with utils.Tempdir() as d:
src = ("""
import other_module
x = other_module.UnusedReferenceNeededToKeepTheImport
class SomeClass(object):
def __init__(will_be_replaced_with_visitor) -> None:
pass
def func(a:SomeClass) -> None:
pass
""")
d.create_file(
"other_module.pyi", """
from typing import Any
def __getattr__(self, name) -> Any: ...""")
ast, loader = self._GetAst(temp_dir=d,
module_name="module1",
src=src)
ast = ast.Visit(RenameVisitor())
pickled_ast_filename = os.path.join(d.path, "module1.pyi.pickled")
module_map = self._StoreAst(d,
"module1",
pickled_ast_filename,
ast=ast,
loader=loader)
del module_map["module1"]
serialized_ast = pytd_utils.LoadPickle(pickled_ast_filename)
loaded_ast = serialize_ast.ProcessAst(serialized_ast, module_map)
# Look up the "SomeClass" in "def func(a: SomeClass), then run
# VerifyLookup on it. We can't run VerifyLookup on the root node, since
# visitors don't descend into the "cls" attribute of ClassType.
cls = loaded_ast.functions[0].signatures[0].params[0].type.cls
cls.Visit(visitors.VerifyLookup())
