Function _get_formatter

python/tvm/script/highlight.py:99–161

(formatter: str | None = None)

@functools.lru_cache
def _get_formatter(formatter: str | None = None):
    def get_ruff_formatter():
        if shutil.which("ruff") is None:
            return None

        def formatter(code_str):
            proc = subprocess.Popen(
                ["ruff", "format", "--stdin-filename=TVMScript"],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                encoding="utf-8",
            )
            stdout, _stderr = proc.communicate(code_str)
            return stdout

        return formatter

    def get_black_formatter():
        try:
            # pylint: disable=import-outside-toplevel
            import black
        except ImportError:
            return None

        def formatter(code_str):
            return black.format_str(code_str, mode=black.FileMode())

        return formatter

    def get_fallback_formatter():
        def formatter(code_str):
            with warnings.catch_warnings():
                warnings.simplefilter("once", UserWarning)
                ruff_install_cmd = sys.executable + " -m pip install ruff"
                black_install_cmd = (
                    sys.executable + ' -m pip install "black==22.3.0" --upgrade --user'
                )
                warnings.warn(
                    f"Neither the 'ruff' formatter nor the 'black' formatter is available.  "
                    f"To print formatted TVM script, please a formatter.  \n"
                    f"To install ruff: {ruff_install_cmd}\n"
                    f"To install black: {black_install_cmd}",
                    category=UserWarning,
                )
            return code_str

        return formatter

    # formatter = "black"
    if formatter is None:
        options = [get_ruff_formatter, get_black_formatter]
    elif formatter == "ruff":
        options = [get_ruff_formatter]
    elif formatter == "black":
        options = [get_black_formatter]
    else:
        raise ValueError(f"Unknown formatter: {formatter}")