pytorch

1
import ast
2
import functools
3
import inspect
4
from textwrap import dedent
5
from typing import Any, List, NamedTuple, Optional, Tuple
6

7
from torch._C import ErrorReport
8
from torch._C._jit_tree_views import SourceRangeFactory
9

10

11
def get_source_lines_and_file(
12
    obj: Any,
13
    error_msg: Optional[str] = None,
14
) -> Tuple[List[str], int, Optional[str]]:
15
    """
16
    Wrapper around inspect.getsourcelines and inspect.getsourcefile.
17

18
    Returns: (sourcelines, file_lino, filename)
19
    """
20
    filename = None  # in case getsourcefile throws
21
    try:
22
        filename = inspect.getsourcefile(obj)
23
        sourcelines, file_lineno = inspect.getsourcelines(obj)
24
    except OSError as e:
25
        msg = (
26
            f"Can't get source for {obj}. TorchScript requires source access in "
27
            "order to carry out compilation, make sure original .py files are "
28
            "available."
29
        )
30
        if error_msg:
31
            msg += "\n" + error_msg
32
        raise OSError(msg) from e
33

34
    return sourcelines, file_lineno, filename
35

36

37
def normalize_source_lines(sourcelines: List[str]) -> List[str]:
38
    """
39
    This helper function accepts a list of source lines. It finds the
40
    indentation level of the function definition (`def`), then it indents
41
    all lines in the function body to a point at or greater than that
42
    level. This allows for comments and continued string literals that
43
    are at a lower indentation than the rest of the code.
44
    Args:
45
        sourcelines: function source code, separated into lines by
46
                        the '\n' character
47
    Returns:
48
        A list of source lines that have been correctly aligned
49
    """
50

51
    def remove_prefix(text, prefix):
52
        return text[text.startswith(prefix) and len(prefix) :]
53

54
    # Find the line and line number containing the function definition
55
    idx = None
56
    for i, l in enumerate(sourcelines):
57
        if l.lstrip().startswith("def"):
58
            idx = i
59
            break
60

61
    # This will happen when the function is a lambda- we won't find "def" anywhere in the source
62
    # lines in that case. Currently trying to JIT compile a lambda will throw an error up in
63
    # `parse_def()`, but we might want to handle this case in the future.
64
    if idx is None:
65
        return sourcelines
66

67
    # Get a string representing the amount of leading whitespace
68
    fn_def = sourcelines[idx]
69
    whitespace = fn_def.split("def")[0]
70

71
    # Add this leading whitespace to all lines before and after the `def`
72
    aligned_prefix = [
73
        whitespace + remove_prefix(s, whitespace) for s in sourcelines[:idx]
74
    ]
75
    aligned_suffix = [
76
        whitespace + remove_prefix(s, whitespace) for s in sourcelines[idx + 1 :]
77
    ]
78

79
    # Put it together again
80
    aligned_prefix.append(fn_def)
81
    return aligned_prefix + aligned_suffix
82

83

84
# Thin wrapper around SourceRangeFactory to store extra metadata
85
# about the function-to-be-compiled.
86
class SourceContext(SourceRangeFactory):
87
    def __init__(
88
        self,
89
        source,
90
        filename,
91
        file_lineno,
92
        leading_whitespace_len,
93
        uses_true_division=True,
94
        funcname=None,
95
    ):
96
        super().__init__(source, filename, file_lineno, leading_whitespace_len)
97
        self.uses_true_division = uses_true_division
98
        self.filename = filename
99
        self.funcname = funcname
100

101

102
@functools.lru_cache(maxsize=None)
103
def make_source_context(*args):
104
    return SourceContext(*args)
105

106

107
def fake_range():
108
    return SourceContext("", None, 0, 0).make_raw_range(0, 1)
109

110

111
class ParsedDef(NamedTuple):
112
    ast: ast.Module
113
    ctx: SourceContext
114
    source: str
115
    filename: Optional[str]
116
    file_lineno: int
117

118

119
def parse_def(fn):
120
    sourcelines, file_lineno, filename = get_source_lines_and_file(
121
        fn, ErrorReport.call_stack()
122
    )
123
    sourcelines = normalize_source_lines(sourcelines)
124
    source = "".join(sourcelines)
125
    dedent_src = dedent(source)
126
    py_ast = ast.parse(dedent_src)
127
    if len(py_ast.body) != 1 or not isinstance(py_ast.body[0], ast.FunctionDef):
128
        raise RuntimeError(
129
            f"Expected a single top-level function: {filename}:{file_lineno}"
130
        )
131
    leading_whitespace_len = len(source.split("\n", 1)[0]) - len(
132
        dedent_src.split("\n", 1)[0]
133
    )
134
    ctx = make_source_context(
135
        source, filename, file_lineno, leading_whitespace_len, True, fn.__name__
136
    )
137
    return ParsedDef(py_ast, ctx, source, filename, file_lineno)
138