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 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 testClearClassPointers(self):
cls = pytd.Class("foo", None, (), (), (), (), None, ())
t = pytd.ClassType("foo", cls)
t = t.Visit(visitors.ClearClassPointers())
self.assertIsNone(t.cls)
