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"))