96 lines
3.9 KiB
Python
96 lines
3.9 KiB
Python
|
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
|
||
|
|
"""
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
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
|