Add ufmt to unify black and usort (#81157)

This reworks [80257](https://github.com/pytorch/pytorch/pull/80257) a bit to use ufmt:

* ufmt https://ufmt.omnilib.dev/en/stable/ unifies both black and usort to automatically format the file in the "most Pythonic" way
* Also make a demo run for all files in `tools/linter/**/*.py`

Pull Request resolved: https://github.com/pytorch/pytorch/pull/81157
Approved by: https://github.com/suo

.lintrunner.toml

@@ -653,3 +653,31 @@ init_command = [
     '--dry-run={{DRYRUN}}',
     'PyYAML==6.0',
 ]
+
+# This also provides the same functionality as BLACK formatter. Just to be
+# on the safe side, we will run both BLACK and UFMT for a while to make sure
+# that nothing breaks before removing the former
+[[linter]]
+code = 'UFMT'
+# This is a big bang change to reformat every python files in the repo,
+# so it's arguably safer to gradually reformat them in subsequent PR.
+# After that, we can use **/*.py to apply the formatter to everything
+include_patterns = ['tools/linter/**/*.py']
+command = [
+    'python3',
+    'tools/linter/adapters/ufmt_linter.py',
+    '--',
+    '@{{PATHSFILE}}'
+]
+exclude_patterns = [
+    'tools/gen_vulkan_spv.py',
+]
+init_command = [
+    'python3',
+    'tools/linter/adapters/pip_init.py',
+    '--dry-run={{DRYRUN}}',
+    'black==22.3.0',
+    'ufmt==1.3.3',
+    'usort==1.0.2',
+]
+is_formatter = true

mypy-strict.ini

@@ -64,3 +64,6 @@ ignore_missing_imports = True
 
 [mypy-mypy.*]
 ignore_missing_imports = True
+
+[mypy-usort.*]
+ignore_missing_imports = True

mypy.ini

@@ -284,3 +284,6 @@ ignore_missing_imports = True
 
 [mypy-dill.*]
 ignore_missing_imports = True
+
+[mypy-usort.*]
+ignore_missing_imports = True

tools/linter/adapters/actionlint_linter.py

@@ -1,9 +1,9 @@
 import argparse
 import concurrent.futures
-import os
-import re
 import json
 import logging
+import os
+import re
 import subprocess
 import time
 from enum import Enum

tools/linter/adapters/black_linter.py

@@ -7,7 +7,7 @@ import subprocess
 import sys
 import time
 from enum import Enum
-from typing import Any, List, NamedTuple, Optional, BinaryIO
+from typing import Any, BinaryIO, List, NamedTuple, Optional
 
 
 IS_WINDOWS: bool = os.name == "nt"

tools/linter/adapters/circleci_linter.py

@@ -2,15 +2,15 @@
 Checks that the configuration in .circleci/config.yml has been properly regenerated.
 """
 
-import os
 import argparse
+import json
+import logging
+import os
 import subprocess
 import sys
-import logging
 import time
 from enum import Enum
 from typing import List, NamedTuple, Optional
-import json
 
 
 CHECKED_IN_FILE = "config.yml"

tools/linter/adapters/flake8_linter.py

@@ -7,7 +7,7 @@ import subprocess
 import sys
 import time
 from enum import Enum
-from typing import Any, Dict, List, NamedTuple, Optional, Set, Pattern
+from typing import Any, Dict, List, NamedTuple, Optional, Pattern, Set
 
 
 IS_WINDOWS: bool = os.name == "nt"

tools/linter/adapters/nativefunctions_linter.py

@@ -14,14 +14,15 @@ is simply to make sure that there is *some* configuration of ruamel that can rou
 the YAML, not to be prescriptive about it.
 """
 
-import ruamel.yaml  # type: ignore[import]
 import argparse
 import json
 import sys
-from io import StringIO
 from enum import Enum
+from io import StringIO
 from typing import NamedTuple, Optional
 
+import ruamel.yaml  # type: ignore[import]
+
 
 class LintSeverity(str, Enum):
     ERROR = "error"

tools/linter/adapters/pip_init.py

@@ -1,9 +1,9 @@
 """
 Initializer script that installs stuff to pip.
 """
-import os
 import argparse
 import logging
+import os
 import subprocess
 import sys
 import time

tools/linter/adapters/shellcheck_linter.py

@@ -1,9 +1,9 @@
 import argparse
 import json
 import logging
+import shutil
 import subprocess
 import time
-import shutil
 from enum import Enum
 from typing import List, NamedTuple, Optional
 

tools/linter/adapters/testowners_linter.py

@@ -8,10 +8,10 @@ has valid ownership information in a comment header. Valid means:
   - Each owner label actually exists in PyTorch
   - Each owner label starts with "module: " or "oncall: " or is in ACCEPTABLE_OWNER_LABELS
 """
-import json
 import argparse
+import json
 from enum import Enum
-from typing import List, Any, Optional, NamedTuple
+from typing import Any, List, NamedTuple, Optional
 from urllib.request import urlopen
 
 

tools/linter/adapters/ufmt_linter.py

@@ -0,0 +1,141 @@
+import argparse
+import concurrent.futures
+import json
+import logging
+import os
+import sys
+from enum import Enum
+from pathlib import Path
+from typing import Any, List, NamedTuple, Optional
+
+from ufmt.core import make_black_config, ufmt_string
+from usort import Config as UsortConfig
+
+
+IS_WINDOWS: bool = os.name == "nt"
+
+
+def eprint(*args: Any, **kwargs: Any) -> None:
+    print(*args, file=sys.stderr, flush=True, **kwargs)
+
+
+class LintSeverity(str, Enum):
+    ERROR = "error"
+    WARNING = "warning"
+    ADVICE = "advice"
+    DISABLED = "disabled"
+
+
+class LintMessage(NamedTuple):
+    path: Optional[str]
+    line: Optional[int]
+    char: Optional[int]
+    code: str
+    severity: LintSeverity
+    name: str
+    original: Optional[str]
+    replacement: Optional[str]
+    description: Optional[str]
+
+
+def as_posix(name: str) -> str:
+    return name.replace("\\", "/") if IS_WINDOWS else name
+
+
+def format_error_message(filename: str, err: Exception) -> LintMessage:
+    return LintMessage(
+        path=filename,
+        line=None,
+        char=None,
+        code="UFMT",
+        severity=LintSeverity.ADVICE,
+        name="command-failed",
+        original=None,
+        replacement=None,
+        description=(f"Failed due to {err.__class__.__name__}:\n{err}"),
+    )
+
+
+def check_file(
+    filename: str,
+) -> List[LintMessage]:
+    with open(filename, "rb") as f:
+        original = f.read().decode("utf-8")
+
+    try:
+        path = Path(filename)
+
+        usort_config = UsortConfig.find(path)
+        black_config = make_black_config(path)
+
+        # Use UFMT API to call both usort and black
+        replacement = ufmt_string(
+            path=path,
+            content=original,
+            usort_config=usort_config,
+            black_config=black_config,
+        )
+
+        if original == replacement:
+            return []
+
+        return [
+            LintMessage(
+                path=filename,
+                line=None,
+                char=None,
+                code="UFMT",
+                severity=LintSeverity.WARNING,
+                name="format",
+                original=original,
+                replacement=replacement,
+                description="Run `lintrunner -a` to apply this patch.",
+            )
+        ]
+    except Exception as err:
+        return [format_error_message(filename, err)]
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Format files with ufmt (black + usort).",
+        fromfile_prefix_chars="@",
+    )
+    parser.add_argument(
+        "--verbose",
+        action="store_true",
+        help="verbose logging",
+    )
+    parser.add_argument(
+        "filenames",
+        nargs="+",
+        help="paths to lint",
+    )
+    args = parser.parse_args()
+
+    logging.basicConfig(
+        format="<%(threadName)s:%(levelname)s> %(message)s",
+        level=logging.NOTSET
+        if args.verbose
+        else logging.DEBUG
+        if len(args.filenames) < 1000
+        else logging.INFO,
+        stream=sys.stderr,
+    )
+
+    with concurrent.futures.ThreadPoolExecutor(
+        max_workers=os.cpu_count(),
+        thread_name_prefix="Thread",
+    ) as executor:
+        futures = {executor.submit(check_file, x): x for x in args.filenames}
+        for future in concurrent.futures.as_completed(futures):
+            try:
+                for lint_message in future.result():
+                    print(json.dumps(lint_message._asdict()), flush=True)
+            except Exception:
+                logging.critical('Failed at "%s".', futures[future])
+                raise
+
+
+if __name__ == "__main__":
+    main()

tools/linter/adapters/update_s3.py

@@ -6,11 +6,12 @@ for how to configure them.
 """
 
 import argparse
-import boto3  # type: ignore[import]
-import json
 import hashlib
-import os
+import json
 import logging
+import os
+
+import boto3  # type: ignore[import]
 
 
 def compute_file_sha256(path: str) -> str:

tools/linter/adapters/workflow_consistency_linter.py

@@ -5,12 +5,12 @@ Any job with a specific `sync-tag` must match all other jobs with the same `sync
 import argparse
 import itertools
 import json
-from pathlib import Path
-from typing import Iterable, Any, Optional, NamedTuple, Dict
-from enum import Enum
 from collections import defaultdict
+from enum import Enum
+from pathlib import Path
+from typing import Any, Dict, Iterable, NamedTuple, Optional
 
-from yaml import load, CSafeLoader, dump
+from yaml import CSafeLoader, dump, load
 
 
 class LintSeverity(str, Enum):

tools/linter/clang_tidy/generate_build_files.py

@@ -1,6 +1,6 @@
+import os
 import subprocess
 import sys
-import os
 from typing import List