def StoreAst(ast, filename):
"""Loads and stores an ast to disk.
Args:
ast: The pytd.TypeDeclUnit to save to disk.
filename: The filename for the pickled output
Returns:
True iff the save operation was successful.
"""
if ast.name.endswith(".__init__"):
ast = ast.Visit(
RenameModuleVisitor(ast.name,
ast.name.rsplit(".__init__", 1)[0]))
# Collect dependencies
deps = visitors.CollectDependencies()
ast.Visit(deps)
dependencies = deps.modules or set()
# Clean external references
ast.Visit(visitors.ClearClassPointers())
indexer = FindClassTypesVisitor()
ast.Visit(indexer)
utils.SavePickle(
SerializableAst(ast, list(sorted(dependencies)),
list(sorted(indexer.class_type_nodes))), filename)
return True
def StoreAst(ast, filename=None):
"""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.
Returns:
The pickled string, if no filename was given. (None otherwise.)
"""
if ast.name.endswith(".__init__"):
ast = ast.Visit(
RenameModuleVisitor(ast.name,
ast.name.rsplit(".__init__", 1)[0]))
# Collect dependencies
deps = visitors.CollectDependencies()
ast.Visit(deps)
dependencies = deps.modules or set()
# Clean external references
ast.Visit(visitors.ClearClassPointers())
indexer = FindClassTypesVisitor()
ast.Visit(indexer)
return pytd_utils.SavePickle(
SerializableAst(ast, list(sorted(dependencies)),
list(sorted(indexer.class_type_nodes))), filename)
def testCollectDependencies(self):
src = textwrap.dedent("""
l = ... # type: list[int or baz.BigInt]
def f1() -> bar.Bar
def f2() -> foo.bar.Baz
""")
deps = visitors.CollectDependencies()
self.Parse(src).Visit(deps)
self.assertSetEqual({"baz", "bar", "foo.bar"}, deps.modules)
def _verifyDeps(self, module, immediate_deps, late_deps):
if isinstance(module, str):
data = cPickle.loads(module)
self.assertItemsEqual(data.dependencies, immediate_deps)
ast = data.ast
else:
c = visitors.CollectDependencies()
module.Visit(c)
self.assertItemsEqual(c.modules, immediate_deps)
ast = module
c = visitors.CollectLateDependencies()
ast.Visit(c)
self.assertItemsEqual(c.modules, late_deps)
def resolve_ast(self, ast):
"""Fill in all ExternalType.cls pointers."""
deps = visitors.CollectDependencies()
ast.Visit(deps)
if deps.modules:
for name in deps.modules:
if name not in self._modules:
self.import_name(name)
self._resolve_all()
module_map = {
name: module.ast
for name, module in self._modules.items()
}
ast.Visit(visitors.InPlaceLookupExternalClasses(module_map, True))
return ast
def _load_and_resolve_ast_dependencies(self, ast, ast_name=None):
"""Fill in all ClassType.cls pointers."""
deps = visitors.CollectDependencies()
ast.Visit(deps)
if deps.modules:
for name in deps.modules:
if name not in self._modules:
other_ast = self._import_name(name)
if other_ast is None:
raise DependencyNotFoundError(name, ast_name
or ast.name)
module_map = {
name: module.ast
for name, module in self._modules.items()
}
ast = ast.Visit(
visitors.LookupExternalTypes(module_map,
full_names=True,
self_name=ast_name))
return ast
def _load_and_resolve_ast_dependencies(self, ast, ast_name=None):
"""Fill in all ClassType.cls pointers."""
deps = visitors.CollectDependencies()
ast.Visit(deps)
if deps.modules:
for name in deps.modules:
if name not in self._modules:
other_ast = self._import_name(name)
if other_ast is None:
error = "Can't find pyi for %r" % name
raise BadDependencyError(error, ast_name or ast.name)
module_map = {
name: module.ast
for name, module in self._modules.items()
}
try:
ast = ast.Visit(
visitors.LookupExternalTypes(module_map,
full_names=True,
self_name=ast_name))
except KeyError as e:
raise BadDependencyError(e.message, ast_name or ast.name)
return ast
def _collect_ast_dependencies(self, ast): """Goes over an ast and returns all references module names.""" deps = visitors.CollectDependencies() ast.Visit(deps) return deps.modules
