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().load_file(module_name, filename, ast) existing = self._get_existing_ast(module_name) if existing: return existing loaded_ast = pytd_utils.LoadPickle(filename, open_function=self.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, 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) 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].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(utils.message(e), module_name) self._modules[module_name].ast = ast self._modules[module_name].pickle = None self._modules[module_name].dirty = False return ast
def _unpickle_module(self, module): if not module.pickle: return todo = [module] seen = set() newly_loaded_asts = [] while todo: m = todo.pop() if m in seen: continue else: seen.add(m) if not m.pickle: continue loaded_ast = cPickle.loads(m.pickle) deps = [ d for d, _ in loaded_ast.dependencies if d != loaded_ast.ast.name ] loaded_ast = serialize_ast.EnsureAstName(loaded_ast, m.module_name) assert m.module_name in self._modules todo.extend(self._modules[dependency] for dependency in deps) newly_loaded_asts.append(loaded_ast) m.ast = loaded_ast.ast m.pickle = None module_map = self._get_module_map() for loaded_ast in newly_loaded_asts: unused_new_serialize_ast = serialize_ast.FillLocalReferences( loaded_ast, module_map) assert module.ast
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 _unpickle_module(self, module): """Unpickle a pickled ast and its dependencies.""" if not module.pickle: return todo = [module] seen = set() newly_loaded_asts = [] while todo: m = todo.pop() if m in seen: continue else: seen.add(m) if not m.pickle: continue loaded_ast = pickle.loads(m.pickle) deps = [ d for d, _ in loaded_ast.dependencies if d != loaded_ast.ast.name ] loaded_ast = serialize_ast.EnsureAstName(loaded_ast, m.module_name) assert m.module_name in self._modules for dependency in deps: module_prefix = dependency while module_prefix == dependency or "." in module_prefix: if module_prefix in self._modules: break module_prefix, _, _ = module_prefix.rpartition(".") else: raise KeyError(f"Module not found: {dependency}") todo.append(self._modules[module_prefix]) newly_loaded_asts.append(loaded_ast) m.ast = loaded_ast.ast m.pickle = None module_map = self.get_module_map() for loaded_ast in newly_loaded_asts: serialize_ast.FillLocalReferences(loaded_ast, module_map) assert module.ast