def assertBuiltinsPickleEqual(self, f1, f2): with open(f1, "rb") as pickle1, open(f2, "rb") as pickle2: if pickle1.read() == pickle2.read(): return out1 = pytd_utils.LoadPickle(f1, compress=True) out2 = pytd_utils.LoadPickle(f2, compress=True) raise AssertionError("\n".join(pytd_utils.DiffNamedPickles(out1, out2)))
def assertBuiltinsPickleEqual(self, f1, f2): with open(f1, "rb") as pickle1, open(f2, "rb") as pickle2: if pickle1.read() == pickle2.read(): return out1 = pytd_utils.LoadPickle(f1, compress=True) out2 = pytd_utils.LoadPickle(f2, compress=True) diff = [] for (name1, ast1), (name2, ast2) in zip(out1, out2): if name1 != name2: raise AssertionError("different ordering of pyi files") elif ast1 != ast2: diff.append(name1) raise AssertionError("pyi files differ: " + ", ".join(diff))
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 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 load_from_pickle(cls, filename, base_module, **kwargs): items = pytd_utils.LoadPickle(filename, compress=True) modules = { name: Module(name, filename=None, ast=None, pickle=pickle, dirty=False) for name, pickle in items } return cls(base_module=base_module, modules=modules, **kwargs)
def load_file(self, module_name, filename, ast=None): """Load (or retrieve from cache) a module and resolve its dependencies.""" if not is_pickle(filename): return super(PickledPyiLoader, self).load_file(module_name, filename, ast) existing = self._get_existing_ast(module_name) if existing: # TODO(kramm): When does this happen? return existing loaded_ast = pytd_utils.LoadPickle(filename) # At this point ast.name and module_name could be different. # They are later synced in ProcessAst. dependencies = {d: names for d, names in loaded_ast.dependencies if d != loaded_ast.ast.name} loaded_ast = serialize_ast.EnsureAstName(loaded_ast, module_name, fix=True) self._modules[module_name] = Module(module_name, filename, loaded_ast.ast) self._load_ast_dependencies(dependencies, ast, module_name) try: ast = serialize_ast.ProcessAst(loaded_ast, self._get_module_map()) except serialize_ast.UnrestorableDependencyError as e: del self._modules[module_name] raise BadDependencyError(utils.message(e), module_name) # Mark all the module's late dependencies as explicitly imported. for d, _ in loaded_ast.late_dependencies: if d != loaded_ast.ast.name: self.add_module_prefixes(d) self._modules[module_name].ast = ast self._modules[module_name].pickle = None self._modules[module_name].dirty = False return ast
def load_file(self, module_name, filename, ast=None): """Load (or retrieve from cache) a module and resolve its dependencies.""" if not os.path.splitext(filename)[1].startswith(".pickled"): return super(PickledPyiLoader, self).load_file(module_name, filename, ast) existing = self._get_existing_ast(module_name) if existing: # TODO(kramm): When does this happen? return existing loaded_ast = pytd_utils.LoadPickle(filename) # At this point ast.name and module_name could be different. # They are later synced in ProcessAst. dependencies = [ d for d in loaded_ast.dependencies if d != loaded_ast.ast.name ] loaded_ast = serialize_ast.EnsureAstName(loaded_ast, module_name, fix=True) self._modules[module_name] = Module(module_name, filename, loaded_ast.ast) self._load_ast_dependencies(dependencies, ast, module_name) try: ast = serialize_ast.ProcessAst(loaded_ast, self._get_module_map()) except serialize_ast.UnrestorableDependencyError as e: del self._modules[module_name] raise BadDependencyError(e.message, module_name) self._modules[module_name].ast = ast self._modules[module_name].pickle = None self._modules[module_name].dirty = False return ast
def LoadPrecompiled(filename, python_version): """Load precompiled builtins from the specified file.""" # TODO(steenbuck): This should check that the python_version in the current # process is the same as the one used to generate the cached file. global _cached_builtins_pytd assert _cached_builtins_pytd.cache is None _cached_builtins_pytd = Cache(python_version, pytd_utils.LoadPickle(filename))
def testUnrestorableDependencyErrorWithModuleIndex(self): with utils.Tempdir() as d: module_name = "module1" pickled_ast_filename = os.path.join(d.path, "module1.pyi.pickled") module_map = self._StoreAst(d, module_name, pickled_ast_filename) module_map = {} # Remove module2 with self.assertRaises(serialize_ast.UnrestorableDependencyError): serialize_ast.ProcessAst( pytd_utils.LoadPickle(pickled_ast_filename), module_map)
def test_function_type(self): with file_utils.Tempdir() as d: foo = d.create_file("foo.pickle") module_map = self._store_ast(d, "foo", foo, ast=self._get_ast(d, "foo")) p = pytd_utils.LoadPickle(foo) self.assertTrue(p.function_type_nodes) ast = serialize_ast.ProcessAst(p, module_map) f, = [a for a in ast.aliases if a.name == "foo.f"] signature, = f.type.function.signatures self.assertIsNotNone(signature.return_type.cls)
def testUnrestorableDependencyErrorWithoutModuleIndex(self): with utils.Tempdir() as d: module_name = "module1" pickled_ast_filename = os.path.join(d.path, "module1.pyi.pickled") module_map = self._StoreAst(d, module_name, pickled_ast_filename) module_map = {} # Remove module2 loaded_ast = pytd_utils.LoadPickle(pickled_ast_filename) loaded_ast.modified_class_types = None # Remove the index with self.assertRaises(serialize_ast.UnrestorableDependencyError): serialize_ast.ProcessAst(loaded_ast, module_map)
def load_from_pickle(cls, filename, base_module, **kwargs): """Load a pytd module from a pickle file.""" items = pytd_utils.LoadPickle( filename, compress=True, open_function=kwargs.get("open_function", open)) modules = { name: Module(name, filename=None, ast=None, pickle=pickle, has_unresolved_pointers=False) for name, pickle in items } return cls(base_module=base_module, modules=modules, **kwargs)
def load_from_pickle(cls, filename, options): """Load a pytd module from a pickle file.""" items = pytd_utils.LoadPickle(filename, compress=True, open_function=options.open_function) modules = { name: Module(name, filename=None, ast=None, pickle=pickle, has_unresolved_pointers=False) for name, pickle in items } return cls(options, modules=modules)
def testStoreRemovesInit(self): with utils.Tempdir() as d: original_module_name = "module1.__init__" pickled_ast_filename = os.path.join(d.path, "module1.pyi.pickled") module_map = self._StoreAst(d, original_module_name, pickled_ast_filename) serializable_ast = pytd_utils.LoadPickle(pickled_ast_filename) expected_name = "module1" # Check that the module had the expected name before. self.assertTrue(original_module_name in module_map) # Check that module1 wasn't created before storing. self.assertTrue(expected_name not in module_map) # Check that the saved ast had its name changed. self.assertEqual(serializable_ast.ast.name, expected_name)
def testLoadTopLevel(self): """Tests that a pickled file can be read.""" with utils.Tempdir() as d: module_name = "module1" pickled_ast_filename = os.path.join(d.path, "module1.pyi.pickled") module_map = self._StoreAst(d, module_name, pickled_ast_filename) original_ast = module_map[module_name] del module_map[module_name] loaded_ast = serialize_ast.ProcessAst( pytd_utils.LoadPickle(pickled_ast_filename), module_map) self.assertTrue(loaded_ast) self.assertTrue(loaded_ast is not original_ast) self.assertEqual(loaded_ast.name, module_name) self.assertTrue(original_ast.ASTeq(loaded_ast)) loaded_ast.Visit(visitors.VerifyLookup())
def testLoadWithDifferentModuleName(self): with utils.Tempdir() as d: original_module_name = "module1" pickled_ast_filename = os.path.join(d.path, "module1.pyi.pickled") module_map = self._StoreAst(d, original_module_name, pickled_ast_filename) original_ast = module_map[original_module_name] del module_map[original_module_name] new_module_name = "wurstbrot.module2" serializable_ast = pytd_utils.LoadPickle(pickled_ast_filename) serializable_ast = serialize_ast.EnsureAstName( serializable_ast, new_module_name) loaded_ast = serialize_ast.ProcessAst(serializable_ast, module_map) self.assertTrue(loaded_ast) self.assertTrue(loaded_ast is not original_ast) self.assertEqual(loaded_ast.name, new_module_name) loaded_ast.Visit(visitors.VerifyLookup()) self.assertFalse(original_ast.ASTeq(loaded_ast)) ast_new_module, _ = self._GetAst(temp_dir=d, module_name=new_module_name) self.assertTrue(ast_new_module.ASTeq(loaded_ast))
def testNodeIndexVisitorUsage(self): """Confirms that the node index is used. This removes the first node from the class_type_nodes list and checks that that node is not updated by ProcessAst. """ with utils.Tempdir() as d: module_name = "module1" pickled_ast_filename = os.path.join(d.path, "module1.pyi.pickled") module_map = self._StoreAst(d, module_name, pickled_ast_filename) del module_map[module_name] serialized_ast = pytd_utils.LoadPickle(pickled_ast_filename) # The sorted makes the testcase more deterministic. serialized_ast = serialized_ast.Replace( class_type_nodes=sorted(serialized_ast.class_type_nodes)[1:]) loaded_ast = serialize_ast.ProcessAst(serialized_ast, module_map) with self.assertRaisesRegexp( ValueError, "Unresolved class: '__builtin__.NoneType'"): loaded_ast.Visit(visitors.VerifyLookup())
def testLoadWithSameModuleName(self): """Explicitly set the module name and reload with the same name. The difference to testLoadTopLevel is that the module name does not match the filelocation. """ with utils.Tempdir() as d: module_name = "foo.bar.module1" pickled_ast_filename = os.path.join(d.path, "module1.pyi.pickled") module_map = self._StoreAst(d, module_name, pickled_ast_filename) original_ast = module_map[module_name] del module_map[module_name] loaded_ast = serialize_ast.ProcessAst( pytd_utils.LoadPickle(pickled_ast_filename), module_map) self.assertTrue(loaded_ast) self.assertTrue(loaded_ast is not original_ast) self.assertEqual(loaded_ast.name, "foo.bar.module1") self.assertTrue(original_ast.ASTeq(loaded_ast)) loaded_ast.Visit(visitors.VerifyLookup())
def load_file(self, module_name, filename, mod_ast=None): """Load (or retrieve from cache) a module and resolve its dependencies.""" if not pytd_utils.IsPickle(filename): return super().load_file(module_name, filename, mod_ast) existing = self._modules.get_existing_ast(module_name) if existing: return existing loaded_ast = pytd_utils.LoadPickle( filename, open_function=self.options.open_function) # At this point ast.name and module_name could be different. # They are later synced in ProcessAst. dependencies = { d: names for d, names in loaded_ast.dependencies if d != loaded_ast.ast.name } loaded_ast = serialize_ast.EnsureAstName(loaded_ast, module_name, fix=True) self._modules[module_name] = Module(module_name, filename, loaded_ast.ast) self._load_ast_dependencies(dependencies, lookup_ast=mod_ast, lookup_ast_name=module_name) try: ast = serialize_ast.ProcessAst(loaded_ast, self._modules.get_module_map()) except serialize_ast.UnrestorableDependencyError as e: del self._modules[module_name] raise BadDependencyError(utils.message(e), module_name) from e # Mark all the module's late dependencies as explicitly imported. for d, _ in loaded_ast.late_dependencies: if d != loaded_ast.ast.name: self.add_module_prefixes(d) self._modules[module_name].ast = ast self._modules[module_name].pickle = None self._modules[module_name].has_unresolved_pointers = False return ast
def testUnrestorableChild(self): # Assume .cls in a ClassType X in module1 was referencing something for # which, Visitors.LookupExternalTypes returned AnythingType. # Now if something in module1 is referencing X external types need to be # resolved before local types, so that we can resolve local types to the # correct ClassType, as the ClassType instance changes, if .cls can not be # filled and instead AnythingType is used. class RenameVisitor(visitors.Visitor): def __init__(self, *args, **kwargs): super(RenameVisitor, self).__init__(*args, **kwargs) self._init = False def EnterFunction(self, func): if func.name == "__init__": self._init = True return None return False def LeaveFunction(self, func): self._init = False def VisitClassType(self, cls_type): if self._init: cls_type = cls_type.Replace( name="other_module.unknown_Reference") # Needs to be copied manually as it is not part of the NamedTuple. cls_type.cls = None return cls_type with utils.Tempdir() as d: src = (""" import other_module x = other_module.UnusedReferenceNeededToKeepTheImport class SomeClass(object): def __init__(will_be_replaced_with_visitor) -> None: pass def func(a:SomeClass) -> None: pass """) d.create_file( "other_module.pyi", """ from typing import Any def __getattr__(self, name) -> Any: ...""") ast, loader = self._GetAst(temp_dir=d, module_name="module1", src=src) ast = ast.Visit(RenameVisitor()) pickled_ast_filename = os.path.join(d.path, "module1.pyi.pickled") module_map = self._StoreAst(d, "module1", pickled_ast_filename, ast=ast, loader=loader) del module_map["module1"] serialized_ast = pytd_utils.LoadPickle(pickled_ast_filename) loaded_ast = serialize_ast.ProcessAst(serialized_ast, module_map) # Look up the "SomeClass" in "def func(a: SomeClass), then run # VerifyLookup on it. We can't run VerifyLookup on the root node, since # visitors don't descend into the "cls" attribute of ClassType. cls = loaded_ast.functions[0].signatures[0].params[0].type.cls cls.Visit(visitors.VerifyLookup())