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 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 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 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 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 Precompile(filename, python_version):
"""Write precompiled builtins to the specified file."""
assert python_version
data = GetBuiltinsAndTyping(python_version)
pytd_utils.SavePickle(data, filename)
