pytorch

1
import argparse
2
import concurrent.futures
3
import json
4
import logging
5
import os
6
import sys
7
from enum import Enum
8
from pathlib import Path
9
from typing import Any, List, NamedTuple, Optional
10

11
from ufmt.core import ufmt_string
12
from ufmt.util import make_black_config
13
from usort import Config as UsortConfig
14

15

16
IS_WINDOWS: bool = os.name == "nt"
17

18

19
def eprint(*args: Any, **kwargs: Any) -> None:
20
    print(*args, file=sys.stderr, flush=True, **kwargs)
21

22

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

29

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

41

42
def as_posix(name: str) -> str:
43
    return name.replace("\\", "/") if IS_WINDOWS else name
44

45

46
def format_error_message(filename: str, err: Exception) -> LintMessage:
47
    return LintMessage(
48
        path=filename,
49
        line=None,
50
        char=None,
51
        code="UFMT",
52
        severity=LintSeverity.ADVICE,
53
        name="command-failed",
54
        original=None,
55
        replacement=None,
56
        description=(f"Failed due to {err.__class__.__name__}:\n{err}"),
57
    )
58

59

60
def check_file(
61
    filename: str,
62
) -> List[LintMessage]:
63
    with open(filename, "rb") as f:
64
        original = f.read().decode("utf-8")
65

66
    try:
67
        path = Path(filename)
68

69
        usort_config = UsortConfig.find(path)
70
        black_config = make_black_config(path)
71

72
        # Use UFMT API to call both usort and black
73
        replacement = ufmt_string(
74
            path=path,
75
            content=original,
76
            usort_config=usort_config,
77
            black_config=black_config,
78
        )
79

80
        if original == replacement:
81
            return []
82

83
        return [
84
            LintMessage(
85
                path=filename,
86
                line=None,
87
                char=None,
88
                code="UFMT",
89
                severity=LintSeverity.WARNING,
90
                name="format",
91
                original=original,
92
                replacement=replacement,
93
                description="Run `lintrunner -a` to apply this patch.",
94
            )
95
        ]
96
    except Exception as err:
97
        return [format_error_message(filename, err)]
98

99

100
def main() -> None:
101
    parser = argparse.ArgumentParser(
102
        description="Format files with ufmt (black + usort).",
103
        fromfile_prefix_chars="@",
104
    )
105
    parser.add_argument(
106
        "--verbose",
107
        action="store_true",
108
        help="verbose logging",
109
    )
110
    parser.add_argument(
111
        "filenames",
112
        nargs="+",
113
        help="paths to lint",
114
    )
115
    args = parser.parse_args()
116

117
    logging.basicConfig(
118
        format="<%(processName)s:%(levelname)s> %(message)s",
119
        level=logging.NOTSET
120
        if args.verbose
121
        else logging.DEBUG
122
        if len(args.filenames) < 1000
123
        else logging.INFO,
124
        stream=sys.stderr,
125
    )
126

127
    with concurrent.futures.ProcessPoolExecutor(
128
        max_workers=os.cpu_count(),
129
    ) as executor:
130
        futures = {executor.submit(check_file, x): x for x in args.filenames}
131
        for future in concurrent.futures.as_completed(futures):
132
            try:
133
                for lint_message in future.result():
134
                    print(json.dumps(lint_message._asdict()), flush=True)
135
            except Exception:
136
                logging.critical('Failed at "%s".', futures[future])
137
                raise
138

139

140
if __name__ == "__main__":
141
    main()
142