pytorch

1
import ast
2
from typing import List, Optional, Tuple
3

4
from ._importlib import _resolve_name
5

6

7
class _ExtractModuleReferences(ast.NodeVisitor):
8
    """
9
    Extract the list of global variables a block of code will read and write
10
    """
11

12
    @classmethod
13
    def run(cls, src: str, package: str) -> List[Tuple[str, Optional[str]]]:
14
        visitor = cls(package)
15
        tree = ast.parse(src)
16
        visitor.visit(tree)
17
        return list(visitor.references.keys())
18

19
    def __init__(self, package):
20
        super().__init__()
21
        self.package = package
22
        self.references = {}
23

24
    def _absmodule(self, module_name: str, level: int) -> str:
25
        if level > 0:
26
            return _resolve_name(module_name, self.package, level)
27
        return module_name
28

29
    def visit_Import(self, node):
30
        for alias in node.names:
31
            self.references[(alias.name, None)] = True
32

33
    def visit_ImportFrom(self, node):
34
        name = self._absmodule(node.module, 0 if node.level is None else node.level)
35
        for alias in node.names:
36
            # from my_package import foo
37
            # foo may be a module, so we have to add it to the list of
38
            # potential references, if import of it fails, we will ignore it
39
            if alias.name != "*":
40
                self.references[(name, alias.name)] = True
41
            else:
42
                self.references[(name, None)] = True
43

44
    def _grab_node_int(self, node):
45
        return node.value
46

47
    def _grab_node_str(self, node):
48
        return node.value
49

50
    def visit_Call(self, node):
51
        # __import__ calls aren't routed to the visit_Import/From nodes
52
        if hasattr(node.func, "id") and node.func.id == "__import__":
53
            try:
54
                name = self._grab_node_str(node.args[0])
55
                fromlist = []
56
                level = 0
57
                if len(node.args) > 3:
58
                    for v in node.args[3].elts:
59
                        fromlist.append(self._grab_node_str(v))
60
                elif hasattr(node, "keywords"):
61
                    for keyword in node.keywords:
62
                        if keyword.arg == "fromlist":
63
                            for v in keyword.value.elts:
64
                                fromlist.append(self._grab_node_str(v))
65
                if len(node.args) > 4:
66
                    level = self._grab_node_int(node.args[4])
67
                elif hasattr(node, "keywords"):
68
                    for keyword in node.keywords:
69
                        if keyword.arg == "level":
70
                            level = self._grab_node_int(keyword.value)
71
                if fromlist == []:
72
                    # the top-level package (the name up till the first dot) is returned
73
                    # when the fromlist argument is empty in normal import system,
74
                    # we need to include top level package to match this behavior and last
75
                    # level package to capture the intended dependency of user
76
                    self.references[(name, None)] = True
77
                    top_name = name.rsplit(".", maxsplit=1)[0]
78
                    if top_name != name:
79
                        top_name = self._absmodule(top_name, level)
80
                        self.references[(top_name, None)] = True
81
                else:
82
                    name = self._absmodule(name, level)
83
                    for alias in fromlist:
84
                        # fromlist args may be submodules, so we have to add the fromlist args
85
                        # to the list of potential references. If import of an arg fails we
86
                        # will ignore it, similar to visit_ImportFrom
87
                        if alias != "*":
88
                            self.references[(name, alias)] = True
89
                        else:
90
                            self.references[(name, None)] = True
91
            except Exception as e:
92
                return
93

94

95
find_files_source_depends_on = _ExtractModuleReferences.run
96