1# mypy: allow-untyped-defs 2import ast 3from typing import List, Optional, Tuple 4 5from ._importlib import _resolve_name 6 7 8class _ExtractModuleReferences(ast.NodeVisitor): 9 """ 10 Extract the list of global variables a block of code will read and write 11 """ 12 13 @classmethod 14 def run(cls, src: str, package: str) -> List[Tuple[str, Optional[str]]]: 15 visitor = cls(package) 16 tree = ast.parse(src) 17 visitor.visit(tree) 18 return list(visitor.references.keys()) 19 20 def __init__(self, package): 21 super().__init__() 22 self.package = package 23 self.references = {} 24 25 def _absmodule(self, module_name: str, level: int) -> str: 26 if level > 0: 27 return _resolve_name(module_name, self.package, level) 28 return module_name 29 30 def visit_Import(self, node): 31 for alias in node.names: 32 self.references[(alias.name, None)] = True 33 34 def visit_ImportFrom(self, node): 35 name = self._absmodule(node.module, 0 if node.level is None else node.level) 36 for alias in node.names: 37 # from my_package import foo 38 # foo may be a module, so we have to add it to the list of 39 # potential references, if import of it fails, we will ignore it 40 if alias.name != "*": 41 self.references[(name, alias.name)] = True 42 else: 43 self.references[(name, None)] = True 44 45 def _grab_node_int(self, node): 46 return node.value 47 48 def _grab_node_str(self, node): 49 return node.value 50 51 def visit_Call(self, node): 52 # __import__ calls aren't routed to the visit_Import/From nodes 53 if hasattr(node.func, "id") and node.func.id == "__import__": 54 try: 55 name = self._grab_node_str(node.args[0]) 56 fromlist = [] 57 level = 0 58 if len(node.args) > 3: 59 for v in node.args[3].elts: 60 fromlist.append(self._grab_node_str(v)) 61 elif hasattr(node, "keywords"): 62 for keyword in node.keywords: 63 if keyword.arg == "fromlist": 64 for v in keyword.value.elts: 65 fromlist.append(self._grab_node_str(v)) 66 if len(node.args) > 4: 67 level = self._grab_node_int(node.args[4]) 68 elif hasattr(node, "keywords"): 69 for keyword in node.keywords: 70 if keyword.arg == "level": 71 level = self._grab_node_int(keyword.value) 72 if fromlist == []: 73 # the top-level package (the name up till the first dot) is returned 74 # when the fromlist argument is empty in normal import system, 75 # we need to include top level package to match this behavior and last 76 # level package to capture the intended dependency of user 77 self.references[(name, None)] = True 78 top_name = name.rsplit(".", maxsplit=1)[0] 79 if top_name != name: 80 top_name = self._absmodule(top_name, level) 81 self.references[(top_name, None)] = True 82 else: 83 name = self._absmodule(name, level) 84 for alias in fromlist: 85 # fromlist args may be submodules, so we have to add the fromlist args 86 # to the list of potential references. If import of an arg fails we 87 # will ignore it, similar to visit_ImportFrom 88 if alias != "*": 89 self.references[(name, alias)] = True 90 else: 91 self.references[(name, None)] = True 92 except Exception as e: 93 return 94 95 96find_files_source_depends_on = _ExtractModuleReferences.run 97