xref: /aosp_15_r20/external/pytorch/tools/linter/adapters/testowners_linter.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1#!/usr/bin/env python3
2"""
3Test ownership was introduced in https://github.com/pytorch/pytorch/issues/66232.
4
5This lint verifies that every Python test file (file that matches test_*.py or *_test.py in the test folder)
6has 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
12from __future__ import annotations
13
14import argparse
15import json
16from enum import Enum
17from typing import Any, NamedTuple
18from urllib.request import urlopen
19
20
21LINTER_CODE = "TESTOWNERS"
22
23
24class LintSeverity(str, Enum):
25    ERROR = "error"
26    WARNING = "warning"
27    ADVICE = "advice"
28    DISABLED = "disabled"
29
30
31class LintMessage(NamedTuple):
32    path: str | None
33    line: int | None
34    char: int | None
35    code: str
36    severity: LintSeverity
37    name: str
38    original: str | None
39    replacement: str | None
40    description: str | None
41
42
43# Team/owner labels usually start with "module: " or "oncall: ", but the following are acceptable exceptions
44ACCEPTABLE_OWNER_LABELS = ["NNC", "high priority"]
45OWNERS_PREFIX = "# Owner(s): "
46
47
48def get_pytorch_labels() -> Any:
49    labels = (
50        urlopen("https://ossci-metrics.s3.amazonaws.com/pytorch_labels.json")
51        .read()
52        .decode("utf-8")
53    )
54    return json.loads(labels)
55
56
57PYTORCH_LABELS = get_pytorch_labels()
58# Team/owner labels usually start with "module: " or "oncall: ", but the following are acceptable exceptions
59ACCEPTABLE_OWNER_LABELS = ["NNC", "high priority"]
60GLOB_EXCEPTIONS = ["**/test/run_test.py"]
61
62
63def check_labels(
64    labels: list[str], filename: str, line_number: int
65) -> list[LintMessage]:
66    lint_messages = []
67    for label in labels:
68        if label not in PYTORCH_LABELS:
69            lint_messages.append(
70                LintMessage(
71                    path=filename,
72                    line=line_number,
73                    char=None,
74                    code=LINTER_CODE,
75                    severity=LintSeverity.ERROR,
76                    name="[invalid-label]",
77                    original=None,
78                    replacement=None,
79                    description=(
80                        f"{label} is not a PyTorch label "
81                        "(please choose from https://github.com/pytorch/pytorch/labels)"
82                    ),
83                )
84            )
85
86        if label.startswith(("module:", "oncall:")) or label in ACCEPTABLE_OWNER_LABELS:
87            continue
88
89        lint_messages.append(
90            LintMessage(
91                path=filename,
92                line=line_number,
93                char=None,
94                code=LINTER_CODE,
95                severity=LintSeverity.ERROR,
96                name="[invalid-owner]",
97                original=None,
98                replacement=None,
99                description=(
100                    f"{label} is not an acceptable owner "
101                    "(please update to another label or edit ACCEPTABLE_OWNERS_LABELS "
102                    "in tools/linters/adapters/testowners_linter.py"
103                ),
104            )
105        )
106
107    return lint_messages
108
109
110def check_file(filename: str) -> list[LintMessage]:
111    lint_messages = []
112    has_ownership_info = False
113
114    with open(filename) as f:
115        for idx, line in enumerate(f):
116            if not line.startswith(OWNERS_PREFIX):
117                continue
118
119            has_ownership_info = True
120            labels = json.loads(line[len(OWNERS_PREFIX) :])
121            lint_messages.extend(check_labels(labels, filename, idx + 1))
122
123    if has_ownership_info is False:
124        lint_messages.append(
125            LintMessage(
126                path=filename,
127                line=None,
128                char=None,
129                code=LINTER_CODE,
130                severity=LintSeverity.ERROR,
131                name="[no-owner-info]",
132                original=None,
133                replacement=None,
134                description="Missing a comment header with ownership information.",
135            )
136        )
137
138    return lint_messages
139
140
141def main() -> None:
142    parser = argparse.ArgumentParser(
143        description="test ownership linter",
144        fromfile_prefix_chars="@",
145    )
146    parser.add_argument(
147        "filenames",
148        nargs="+",
149        help="paths to lint",
150    )
151
152    args = parser.parse_args()
153    lint_messages = []
154
155    for filename in args.filenames:
156        lint_messages.extend(check_file(filename))
157
158    for lint_message in lint_messages:
159        print(json.dumps(lint_message._asdict()), flush=True)
160
161
162if __name__ == "__main__":
163    main()
164