Add usort linter to lintrunner

Summary: tsia

Differential Revision: D36881903

fbshipit-source-id: e0e7675bf7748f0ca138c1993e23cc5954d40430

.lintrunner.toml

@@ -13,3 +13,21 @@ init_command = [
     '--dry-run={{DRYRUN}}',
     'black==22.3.0',
 ]
+is_formatter = true
+
+[[linter]]
+code = 'USORT'
+include_patterns = ['**/*.py']
+command = [
+    'python3',
+    'tools/lint/usort_linter.py',
+    '--',
+    '@{{PATHSFILE}}'
+]
+init_command = [
+    'python3',
+    'tools/lint/pip_init.py',
+    '--dry-run={{DRYRUN}}',
+    'usort==1.0.2',
+]
+is_formatter = true

requirements.txt

@@ -38,4 +38,5 @@ tqdm
 typing-inspect
 typing_extensions
 urllib3
+usort
 websocket-client

tools/lint/usort_linter.py

@@ -0,0 +1,187 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+import concurrent.futures
+import json
+import logging
+import os
+import subprocess
+import sys
+from enum import Enum
+from typing import List, NamedTuple, Optional
+
+from usort import config as usort_config, usort
+
+
+IS_WINDOWS: bool = os.name == "nt"
+
+
+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 check_file(
+    filename: str,
+) -> List[LintMessage]:
+    try:
+        top_of_file_cat = usort_config.Category("top_of_file")
+        known = usort_config.known_factory()
+        # cinder magic imports must be on top (after future imports)
+        known["__strict__"] = top_of_file_cat
+        known["__static__"] = top_of_file_cat
+
+        config = usort_config.Config(
+            categories=(
+                (
+                    usort_config.CAT_FUTURE,
+                    top_of_file_cat,
+                    usort_config.CAT_STANDARD_LIBRARY,
+                    usort_config.CAT_THIRD_PARTY,
+                    usort_config.CAT_FIRST_PARTY,
+                )
+            ),
+            known=known,
+        )
+
+        with open(filename, mode="rb", encoding="utf-8") as f:
+            original = f.read()
+        with open(filename, mode="rb", encoding="utf-8") as f:
+            result = usort(f, config)
+            if result.error:
+                raise result.error
+
+    except subprocess.TimeoutExpired:
+        return [
+            LintMessage(
+                path=filename,
+                line=None,
+                char=None,
+                code="USORT",
+                severity=LintSeverity.ERROR,
+                name="timeout",
+                original=None,
+                replacement=None,
+                description=(
+                    "black timed out while trying to process a file. "
+                    "Please report an issue in pytorch/torchrec."
+                ),
+            )
+        ]
+    except (OSError, subprocess.CalledProcessError) as err:
+        return [
+            LintMessage(
+                path=filename,
+                line=None,
+                char=None,
+                code="USORT",
+                severity=LintSeverity.ADVICE,
+                name="command-failed",
+                original=None,
+                replacement=None,
+                description=(
+                    f"Failed due to {err.__class__.__name__}:\n{err}"
+                    if not isinstance(err, subprocess.CalledProcessError)
+                    else (
+                        "COMMAND (exit code {returncode})\n"
+                        "{command}\n\n"
+                        "STDERR\n{stderr}\n\n"
+                        "STDOUT\n{stdout}"
+                    ).format(
+                        returncode=err.returncode,
+                        command=" ".join(as_posix(x) for x in err.cmd),
+                        stderr=err.stderr.decode("utf-8").strip() or "(empty)",
+                        stdout=err.stdout.decode("utf-8").strip() or "(empty)",
+                    )
+                ),
+            )
+        ]
+
+    replacement = result.output.decode("utf-8")
+    if original == replacement:
+        return []
+
+    return [
+        LintMessage(
+            path=filename,
+            line=None,
+            char=None,
+            code="USORT",
+            severity=LintSeverity.WARNING,
+            name="format",
+            original=original.decode("utf-8"),
+            replacement=replacement.decode("utf-8"),
+            description="Run `lintrunner -a` to apply this patch.",
+        )
+    ]
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Format files with 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, filename): filename
+            for filename 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()