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
