def save_to_pickle(self, filename):
"""Save to a pickle. See PickledPyiLoader.load_from_pickle for reverse."""
# We assume that the Loader is in a consistent state here. In particular, we
# assume that for every module in _modules, all the transitive dependencies
# have been loaded.
items = tuple((name, serialize_ast.StoreAst(module.ast))
for name, module in sorted(self._modules.items()))
# Now pickle the pickles. We keep the "inner" modules as pickles as a
# performance optimization - unpickling is slow.
pytd_utils.SavePickle(items, filename, compress=True)
def test_diff_same_pickle(self):
ast = pytd.TypeDeclUnit("foo", (), (), (), (), ())
with file_utils.Tempdir() as d:
filename = os.path.join(d.path, "foo.pickled")
serialize_ast.StoreAst(ast, filename)
with open(filename, "rb") as fi:
data = fi.read()
named_pickles = [("foo", data)]
self.assertFalse(
pytd_utils.DiffNamedPickles(named_pickles, named_pickles))
def testDiffPickleLength(self):
ast = pytd.TypeDeclUnit("foo", (), (), (), (), ())
with file_utils.Tempdir() as d:
filename = os.path.join(d.path, "foo.pickled")
serialize_ast.StoreAst(ast, filename)
with open(filename, "rb") as fi:
data = fi.read()
named_pickles1 = []
named_pickles2 = [("foo", data)]
self.assertTrue(pytd_utils.DiffNamedPickles(named_pickles1, named_pickles2))
def save_to_pickle(self, filename):
"""Save to a pickle. See PickledPyiLoader.load_from_pickle for reverse."""
# We assume that the Loader is in a consistent state here. In particular, we
# assume that for every module in _modules, all the transitive dependencies
# have been loaded.
items = tuple((name, serialize_ast.StoreAst(module.ast))
for name, module in sorted(self._modules.items()))
# Preparing an ast for pickling clears its class pointers, making it
# unsuitable for reuse, so we have to discard the builtins cache.
builtins.InvalidateCache(self.python_version)
# Now pickle the pickles. We keep the "inner" modules as pickles as a
# performance optimization - unpickling is slow.
pytd_utils.SavePickle(items, filename, compress=True)
def testPickle(self):
with utils.Tempdir() as d:
ast, _ = self._GetAst(temp_dir=d, module_name="foo.bar.module1")
pickled_ast_filename = os.path.join(d.path, "module1.pyi.pickled")
result = serialize_ast.StoreAst(ast, pickled_ast_filename)
self.assertTrue(result)
with open(pickled_ast_filename, "rb") as fi:
serialized_ast = pickle.load(fi)
self.assertTrue(serialized_ast.ast)
self.assertEqual(serialized_ast.dependencies,
["__builtin__", "foo.bar.module1", "module2"])
def testDiffPickleAst(self):
ast1 = pytd.TypeDeclUnit("foo", (), (), (), (), ())
ast2 = ast1.Replace(type_params=(pytd.TypeParameter("T", (), None, None),))
with file_utils.Tempdir() as d:
data = []
for ast in (ast1, ast2):
filename = os.path.join(d.path, "foo.pickled")
serialize_ast.StoreAst(ast, filename)
with open(filename, "rb") as fi:
data.append(fi.read())
named_pickles1 = [("foo", data[0])]
named_pickles2 = [("foo", data[1])]
self.assertTrue(pytd_utils.DiffNamedPickles(named_pickles1, named_pickles2))
def test_pickle(self):
with file_utils.Tempdir() as d:
ast, _ = self._get_ast(temp_dir=d, module_name="foo.bar.module1")
pickled_ast_filename = os.path.join(d.path, "module1.pyi.pickled")
result = serialize_ast.StoreAst(ast, pickled_ast_filename)
self.assertIsNone(result)
with open(pickled_ast_filename, "rb") as fi:
serialized_ast = pickle.load(fi)
self.assertTrue(serialized_ast.ast)
self.assertCountEqual(
dict(serialized_ast.dependencies),
["builtins", "foo.bar.module1", "module2", "queue"])
def _StoreAst(self,
temp_dir,
module_name,
pickled_ast_filename,
ast=None,
loader=None):
if not (ast and loader):
ast, loader = self._GetAst(temp_dir=temp_dir,
module_name=module_name)
serialize_ast.StoreAst(ast, pickled_ast_filename)
module_map = {
name: module.ast
for name, module in loader._modules.items()
}
return module_map
def testLoadWithSameModuleName(self):
with utils.Tempdir() as d:
self._CreateFiles(tempdir=d)
module1 = _Module(module_name="foo.bar.module1", file_name="module1.pyi")
module2 = _Module(module_name="module2", file_name="module2.pyi")
loader, ast = self._LoadAst(tempdir=d, module=module1)
self._PickleModules(loader, d, module1, module2)
pickled_ast_filename = self._GetPath(d, module1.file_name + ".pickled")
result = serialize_ast.StoreAst(ast, pickled_ast_filename)
self.assertIsNone(result)
loaded_ast = self._LoadPickledModule(d, module1)
self.assertTrue(loaded_ast)
self.assertIsNot(loaded_ast, ast)
self.assertTrue(ast.ASTeq(loaded_ast))
loaded_ast.Visit(visitors.VerifyLookup())
def test_load_with_same_module_name(self):
with file_utils.Tempdir() as d:
self._create_files(tempdir=d)
module1 = _Module(module_name="foo.bar.module1", file_name="module1.pyi")
module2 = _Module(module_name="module2", file_name="module2.pyi")
loader, ast = self._load_ast(tempdir=d, module=module1)
self._pickle_modules(loader, d, module1, module2)
pickled_ast_filename = self._get_path(d, module1.file_name + ".pickled")
result = serialize_ast.StoreAst(ast, pickled_ast_filename)
self.assertIsNone(result)
loaded_ast = self._load_pickled_module(d, module1)
self.assertTrue(loaded_ast)
self.assertIsNot(loaded_ast, ast)
self.assertTrue(pytd_utils.ASTeq(ast, loaded_ast))
loaded_ast.Visit(visitors.VerifyLookup())
def write_pickle(ast, options, loader=None):
"""Dump a pickle of the ast to a file."""
loader = loader or load_pytd.create_loader(options)
try:
ast = serialize_ast.PrepareForExport(options.module_name, ast, loader)
except parser.ParseError as e:
if options.nofail:
ast = serialize_ast.PrepareForExport(
options.module_name,
pytd_builtins.GetDefaultAst(options.python_version), loader)
log.warning("***Caught exception: %s", str(e), exc_info=True)
else:
raise
if options.verify_pickle:
ast1 = ast.Visit(visitors.LateTypeToClassType())
ast1 = ast1.Visit(visitors.ClearClassPointers())
ast2 = loader.load_file(options.module_name, options.verify_pickle)
ast2 = ast2.Visit(visitors.ClearClassPointers())
if not pytd_utils.ASTeq(ast1, ast2):
raise AssertionError()
serialize_ast.StoreAst(ast, options.output, options.open_function)
def _Pickle(self, ast, module_name):
assert module_name
ast = serialize_ast.PrepareForExport(module_name, ast, self.loader)
return serialize_ast.StoreAst(ast)
def _pickle_modules(self, loader, tempdir, *modules):
for module in modules:
serialize_ast.StoreAst(
loader._modules[module.module_name].ast,
self._get_path(tempdir, module.file_name + ".pickled"))
def _PickleModules(self, tempdir, *modules):
for module in modules:
serialize_ast.StoreAst(
self.loader._modules[module.module_name].ast,
self._GetPath(tempdir, module.file_name + ".pickled"))
