Spaces:
Running
Running
| import ast | |
| from typing import List, Optional, Tuple | |
| from ._importlib import _resolve_name | |
| class _ExtractModuleReferences(ast.NodeVisitor): | |
| """ | |
| Extract the list of global variables a block of code will read and write | |
| """ | |
| def run(cls, src: str, package: str) -> List[Tuple[str, Optional[str]]]: | |
| visitor = cls(package) | |
| tree = ast.parse(src) | |
| visitor.visit(tree) | |
| return list(visitor.references.keys()) | |
| def __init__(self, package): | |
| super().__init__() | |
| self.package = package | |
| self.references = {} | |
| def _absmodule(self, module_name: str, level: int) -> str: | |
| if level > 0: | |
| return _resolve_name(module_name, self.package, level) | |
| return module_name | |
| def visit_Import(self, node): | |
| for alias in node.names: | |
| self.references[(alias.name, None)] = True | |
| def visit_ImportFrom(self, node): | |
| name = self._absmodule(node.module, 0 if node.level is None else node.level) | |
| for alias in node.names: | |
| # from my_package import foo | |
| # foo may be a module, so we have to add it to the list of | |
| # potential references, if import of it fails, we will ignore it | |
| if alias.name != "*": | |
| self.references[(name, alias.name)] = True | |
| else: | |
| self.references[(name, None)] = True | |
| def _grab_node_int(self, node): | |
| return node.value | |
| def _grab_node_str(self, node): | |
| return node.value | |
| def visit_Call(self, node): | |
| # __import__ calls aren't routed to the visit_Import/From nodes | |
| if hasattr(node.func, "id") and node.func.id == "__import__": | |
| try: | |
| name = self._grab_node_str(node.args[0]) | |
| fromlist = [] | |
| level = 0 | |
| if len(node.args) > 3: | |
| for v in node.args[3].elts: | |
| fromlist.append(self._grab_node_str(v)) | |
| elif hasattr(node, "keywords"): | |
| for keyword in node.keywords: | |
| if keyword.arg == "fromlist": | |
| for v in keyword.value.elts: | |
| fromlist.append(self._grab_node_str(v)) | |
| if len(node.args) > 4: | |
| level = self._grab_node_int(node.args[4]) | |
| elif hasattr(node, "keywords"): | |
| for keyword in node.keywords: | |
| if keyword.arg == "level": | |
| level = self._grab_node_int(keyword.value) | |
| if fromlist == []: | |
| # the top-level package (the name up till the first dot) is returned | |
| # when the fromlist argument is empty in normal import system, | |
| # we need to include top level package to match this behavior and last | |
| # level package to capture the intended dependency of user | |
| self.references[(name, None)] = True | |
| top_name = name.rsplit(".", maxsplit=1)[0] | |
| if top_name != name: | |
| top_name = self._absmodule(top_name, level) | |
| self.references[(top_name, None)] = True | |
| else: | |
| name = self._absmodule(name, level) | |
| for alias in fromlist: | |
| # fromlist args may be submodules, so we have to add the fromlist args | |
| # to the list of potential references. If import of an arg fails we | |
| # will ignore it, similar to visit_ImportFrom | |
| if alias != "*": | |
| self.references[(name, alias)] = True | |
| else: | |
| self.references[(name, None)] = True | |
| except Exception as e: | |
| return | |
| find_files_source_depends_on = _ExtractModuleReferences.run | |