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