pytorch

1
#!/usr/bin/env python3
2
"""
3
Test ownership was introduced in https://github.com/pytorch/pytorch/issues/66232.
4

5
This lint verifies that every Python test file (file that matches test_*.py or *_test.py in the test folder)
6
has valid ownership information in a comment header. Valid means:
7
  - The format of the header follows the pattern "# Owner(s): ["list", "of owner", "labels"]
8
  - Each owner label actually exists in PyTorch
9
  - Each owner label starts with "module: " or "oncall: " or is in ACCEPTABLE_OWNER_LABELS
10
"""
11
import argparse
12
import json
13
from enum import Enum
14
from typing import Any, List, NamedTuple, Optional
15
from urllib.request import urlopen
16

17

18
LINTER_CODE = "TESTOWNERS"
19

20

21
class LintSeverity(str, Enum):
22
    ERROR = "error"
23
    WARNING = "warning"
24
    ADVICE = "advice"
25
    DISABLED = "disabled"
26

27

28
class LintMessage(NamedTuple):
29
    path: Optional[str]
30
    line: Optional[int]
31
    char: Optional[int]
32
    code: str
33
    severity: LintSeverity
34
    name: str
35
    original: Optional[str]
36
    replacement: Optional[str]
37
    description: Optional[str]
38

39

40
# Team/owner labels usually start with "module: " or "oncall: ", but the following are acceptable exceptions
41
ACCEPTABLE_OWNER_LABELS = ["NNC", "high priority"]
42
OWNERS_PREFIX = "# Owner(s): "
43

44

45
def get_pytorch_labels() -> Any:
46
    labels = (
47
        urlopen("https://ossci-metrics.s3.amazonaws.com/pytorch_labels.json")
48
        .read()
49
        .decode("utf-8")
50
    )
51
    return json.loads(labels)
52

53

54
PYTORCH_LABELS = get_pytorch_labels()
55
# Team/owner labels usually start with "module: " or "oncall: ", but the following are acceptable exceptions
56
ACCEPTABLE_OWNER_LABELS = ["NNC", "high priority"]
57
GLOB_EXCEPTIONS = ["**/test/run_test.py"]
58

59

60
def check_labels(
61
    labels: List[str], filename: str, line_number: int
62
) -> List[LintMessage]:
63
    lint_messages = []
64
    for label in labels:
65
        if label not in PYTORCH_LABELS:
66
            lint_messages.append(
67
                LintMessage(
68
                    path=filename,
69
                    line=line_number,
70
                    char=None,
71
                    code=LINTER_CODE,
72
                    severity=LintSeverity.ERROR,
73
                    name="[invalid-label]",
74
                    original=None,
75
                    replacement=None,
76
                    description=(
77
                        f"{label} is not a PyTorch label "
78
                        "(please choose from https://github.com/pytorch/pytorch/labels)"
79
                    ),
80
                )
81
            )
82

83
        if label.startswith(("module:", "oncall:")) or label in ACCEPTABLE_OWNER_LABELS:
84
            continue
85

86
        lint_messages.append(
87
            LintMessage(
88
                path=filename,
89
                line=line_number,
90
                char=None,
91
                code=LINTER_CODE,
92
                severity=LintSeverity.ERROR,
93
                name="[invalid-owner]",
94
                original=None,
95
                replacement=None,
96
                description=(
97
                    f"{label} is not an acceptable owner "
98
                    "(please update to another label or edit ACCEPTABLE_OWNERS_LABELS "
99
                    "in tools/linters/adapters/testowners_linter.py"
100
                ),
101
            )
102
        )
103

104
    return lint_messages
105

106

107
def check_file(filename: str) -> List[LintMessage]:
108
    lint_messages = []
109
    has_ownership_info = False
110

111
    with open(filename) as f:
112
        for idx, line in enumerate(f):
113
            if not line.startswith(OWNERS_PREFIX):
114
                continue
115

116
            has_ownership_info = True
117
            labels = json.loads(line[len(OWNERS_PREFIX) :])
118
            lint_messages.extend(check_labels(labels, filename, idx + 1))
119

120
    if has_ownership_info is False:
121
        lint_messages.append(
122
            LintMessage(
123
                path=filename,
124
                line=None,
125
                char=None,
126
                code=LINTER_CODE,
127
                severity=LintSeverity.ERROR,
128
                name="[no-owner-info]",
129
                original=None,
130
                replacement=None,
131
                description="Missing a comment header with ownership information.",
132
            )
133
        )
134

135
    return lint_messages
136

137

138
def main() -> None:
139
    parser = argparse.ArgumentParser(
140
        description="test ownership linter",
141
        fromfile_prefix_chars="@",
142
    )
143
    parser.add_argument(
144
        "filenames",
145
        nargs="+",
146
        help="paths to lint",
147
    )
148

149
    args = parser.parse_args()
150
    lint_messages = []
151

152
    for filename in args.filenames:
153
        lint_messages.extend(check_file(filename))
154

155
    for lint_message in lint_messages:
156
        print(json.dumps(lint_message._asdict()), flush=True)
157

158

159
if __name__ == "__main__":
160
    main()
161