pytorch

1
#!/usr/bin/env python3
2
"""
3
This lint verifies that every Python test file (file that matches test_*.py or
4
*_test.py in the test folder) has a main block which raises an exception or
5
calls run_tests to ensure that the test will be run in OSS CI.
6

7
Takes ~2 minuters to run without the multiprocessing, probably overkill.
8
"""
9
import argparse
10
import json
11
import multiprocessing as mp
12
from enum import Enum
13
from typing import List, NamedTuple, Optional
14

15
import libcst as cst
16
import libcst.matchers as m
17

18
LINTER_CODE = "TEST_HAS_MAIN"
19

20

21
class HasMainVisiter(cst.CSTVisitor):
22
    def __init__(self) -> None:
23
        super().__init__()
24
        self.found = False
25

26
    def visit_Module(self, node: cst.Module) -> bool:
27
        name = m.Name("__name__")
28
        main = m.SimpleString('"__main__"') | m.SimpleString("'__main__'")
29
        run_test_call = m.Call(
30
            func=m.Name("run_tests") | m.Attribute(attr=m.Name("run_tests"))
31
        )
32
        raise_block = m.Raise()
33

34
        # name == main or main == name
35
        if_main1 = m.Comparison(
36
            name,
37
            [m.ComparisonTarget(m.Equal(), main)],
38
        )
39
        if_main2 = m.Comparison(
40
            main,
41
            [m.ComparisonTarget(m.Equal(), name)],
42
        )
43
        for child in node.children:
44
            if m.matches(child, m.If(test=if_main1 | if_main2)):
45
                if m.findall(child, raise_block | run_test_call):
46
                    self.found = True
47
                    break
48

49
        return False
50

51

52
class LintSeverity(str, Enum):
53
    ERROR = "error"
54
    WARNING = "warning"
55
    ADVICE = "advice"
56
    DISABLED = "disabled"
57

58

59
class LintMessage(NamedTuple):
60
    path: Optional[str]
61
    line: Optional[int]
62
    char: Optional[int]
63
    code: str
64
    severity: LintSeverity
65
    name: str
66
    original: Optional[str]
67
    replacement: Optional[str]
68
    description: Optional[str]
69

70

71
def check_file(filename: str) -> List[LintMessage]:
72
    lint_messages = []
73

74
    with open(filename) as f:
75
        file = f.read()
76
        v = HasMainVisiter()
77
        cst.parse_module(file).visit(v)
78
        if not v.found:
79
            message = (
80
                "Test files need to have a main block which either calls run_tests "
81
                + "(to ensure that the tests are run during OSS CI) or raises an exception "
82
                + "and added to the blocklist in test/run_test.py"
83
            )
84
            lint_messages.append(
85
                LintMessage(
86
                    path=filename,
87
                    line=None,
88
                    char=None,
89
                    code=LINTER_CODE,
90
                    severity=LintSeverity.ERROR,
91
                    name="[no-main]",
92
                    original=None,
93
                    replacement=None,
94
                    description=message,
95
                )
96
            )
97
    return lint_messages
98

99

100
def main() -> None:
101
    parser = argparse.ArgumentParser(
102
        description="test files should have main block linter",
103
        fromfile_prefix_chars="@",
104
    )
105
    parser.add_argument(
106
        "filenames",
107
        nargs="+",
108
        help="paths to lint",
109
    )
110

111
    args = parser.parse_args()
112

113
    pool = mp.Pool(8)
114
    lint_messages = pool.map(check_file, args.filenames)
115
    pool.close()
116
    pool.join()
117

118
    flat_lint_messages = []
119
    for sublist in lint_messages:
120
        flat_lint_messages.extend(sublist)
121

122
    for lint_message in flat_lint_messages:
123
        print(json.dumps(lint_message._asdict()), flush=True)
124

125

126
if __name__ == "__main__":
127
    main()
128