black_isort_linter.py

"""Format file with black and isort."""

# PyTorch LICENSE. See LICENSE file in the root directory of this source tree.

from __future__ import annotations

import argparse
import concurrent.futures
import logging
import os
import subprocess
import sys

import lintrunner_adapters
from lintrunner_adapters import LintMessage, LintSeverity, as_posix, run_command

LINTER_CODE = "BLACK-ISORT"


def check_file(
    filename: str,
    retries: int,
    timeout: int,
    *,
    fast: bool = False,
) -> list[LintMessage]:
    try:
        with open(filename, "rb") as f:
            original = f.read()
        with open(filename, "rb") as f:
            # Run isort first then black so we get consistent result
            # even if isort is not using the black profile
            proc = run_command(
                [sys.executable, "-misort", "-"],
                stdin=f,
                retries=retries,
                timeout=timeout,
                check=True,
            )
            import_sorted = proc.stdout
            # Pipe isort's result to black
            proc = run_command(
                [
                    sys.executable,
                    "-mblack",
                    *(("--pyi",) if filename.endswith(".pyi") else ()),
                    *(("--ipynb",) if filename.endswith(".ipynb") else ()),
                    *(("--fast",) if fast else ()),
                    "--stdin-filename",
                    filename,
                    "-",
                ],
                stdin=None,
                input=import_sorted,
                retries=retries,
                timeout=timeout,
                check=True,
            )
    except subprocess.TimeoutExpired:
        return [
            LintMessage(
                path=filename,
                line=None,
                char=None,
                code=LINTER_CODE,
                severity=LintSeverity.ERROR,
                name="timeout",
                original=None,
                replacement=None,
                description=(
                    "black-isort timed out while trying to process a file. "
                    "Please report an issue in pytorch/pytorch with the "
                    "label 'module: lint'"
                ),
            )
        ]
    except (OSError, subprocess.CalledProcessError) as err:
        return [
            LintMessage(
                path=filename,
                line=None,
                char=None,
                code=LINTER_CODE,
                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 = proc.stdout
    if original == replacement:
        return []

    return [
        LintMessage(
            path=filename,
            line=None,
            char=None,
            code=LINTER_CODE,
            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=f"Format files with black-isort. Linter code: {LINTER_CODE}",
        fromfile_prefix_chars="@",
    )
    parser.add_argument(
        "--fast",
        action="store_true",
        help="If --fast given, skip temporary sanity checks.",
    )
    parser.add_argument(
        "--timeout",
        default=90,
        type=int,
        help="seconds to wait for black-isort",
    )
    lintrunner_adapters.add_default_options(parser)
    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, args.retries, args.timeout): x
            for x in args.filenames
        }
        for future in concurrent.futures.as_completed(futures):
            try:
                for lint_message in future.result():
                    lint_message.display()
            except Exception:
                logging.critical('Failed at "%s".', futures[future])
                raise


if __name__ == "__main__":
    main()