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)