ruff/crates/ty_python_semantic/mdtest.py

"""A runner for Markdown-based tests for ty"""
# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "rich",
#     "watchfiles",
# ]
# ///

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import threading
from pathlib import Path
from typing import Final, Literal, assert_never

from rich.console import Console
from watchfiles import Change, watch

CRATE_NAME: Final = "ty_python_semantic"
CRATE_ROOT: Final = Path(__file__).resolve().parent
TY_VENDORED: Final = CRATE_ROOT.parent / "ty_vendored"
DIRS_TO_WATCH: Final = (
    CRATE_ROOT,
    TY_VENDORED,
    CRATE_ROOT.parent / "ty_test/src",
)
MDTEST_DIR: Final = CRATE_ROOT / "resources" / "mdtest"
MDTEST_README: Final = CRATE_ROOT / "resources" / "README.md"


class MDTestRunner:
    mdtest_executable: Path | None
    console: Console
    filters: list[str]
    enable_external: bool
    upgrade_lockfiles: bool

    def __init__(
        self,
        filters: list[str] | None,
        enable_external: bool,
        upgrade_lockfiles: bool,
    ) -> None:
        self.mdtest_executable = None
        self.console = Console()
        self.filters = [
            f.removesuffix(".md").replace("/", "_").replace("-", "_")
            for f in (filters or [])
        ]
        self.enable_external = enable_external
        self.upgrade_lockfiles = upgrade_lockfiles

    def _run_cargo_test(self, *, message_format: Literal["human", "json"]) -> str:
        return subprocess.check_output(
            [
                "cargo",
                "test",
                "--package",
                CRATE_NAME,
                "--no-run",
                "--color=always",
                "--test=mdtest",
                "--message-format",
                message_format,
            ],
            cwd=CRATE_ROOT,
            env=dict(os.environ, CLI_COLOR="1"),
            stderr=subprocess.STDOUT,
            text=True,
        )

    def _recompile_tests(
        self, status_message: str, *, message_on_success: bool = True
    ) -> bool:
        with self.console.status(status_message):
            # Run it with 'human' format in case there are errors:
            try:
                self._run_cargo_test(message_format="human")
            except subprocess.CalledProcessError as e:
                print(e.output)
                return False

            # Run it again with 'json' format to find the mdtest executable:
            try:
                json_output = self._run_cargo_test(message_format="json")
            except subprocess.CalledProcessError as _:
                # `cargo test` can still fail if something changed in between the two runs.
                # Here we don't have a human-readable output, so just show a generic message:
                self.console.print("[red]Error[/red]: Failed to compile tests")
                return False

            if json_output:
                self._get_executable_path_from_json(json_output)

        if message_on_success:
            self.console.print("[dim]Tests compiled successfully[/dim]")
        return True

    def _get_executable_path_from_json(self, json_output: str) -> None:
        for json_line in json_output.splitlines():
            try:
                data = json.loads(json_line)
            except json.JSONDecodeError:
                continue
            if data.get("target", {}).get("name") == "mdtest":
                self.mdtest_executable = Path(data["executable"])
                break
        else:
            raise RuntimeError(
                "Could not find mdtest executable after successful compilation"
            )

    def _run_mdtest(
        self, arguments: list[str] | None = None, *, capture_output: bool = False
    ) -> subprocess.CompletedProcess:
        assert self.mdtest_executable is not None

        arguments = arguments or []
        return subprocess.run(
            [self.mdtest_executable, *arguments],
            cwd=CRATE_ROOT,
            env=dict(
                os.environ,
                CLICOLOR_FORCE="1",
                INSTA_FORCE_PASS="1",
                INSTA_OUTPUT="none",
                MDTEST_EXTERNAL="1" if self.enable_external else "0",
                MDTEST_UPGRADE_LOCKFILES="1" if self.upgrade_lockfiles else "0",
            ),
            capture_output=capture_output,
            text=True,
            check=False,
        )

    def _run_mdtests_for_file(self, markdown_file: Path) -> None:
        test_name = f"mdtest::{markdown_file}"

        output = self._run_mdtest(["--exact", test_name], capture_output=True)

        if output.returncode == 0:
            if "running 0 tests\n" in output.stdout:
                self.console.log(
                    f"[yellow]Warning[/yellow]: No tests were executed with filter '{test_name}'"
                )
            else:
                self.console.print(
                    f"Test for [bold green]{markdown_file}[/bold green] succeeded"
                )
        else:
            self.console.print()
            self.console.rule(
                f"Test for [bold red]{markdown_file}[/bold red] failed",
                style="gray",
            )
            self._print_trimmed_cargo_test_output(
                output.stdout + output.stderr, test_name
            )

    def _print_trimmed_cargo_test_output(self, output: str, test_name: str) -> None:
        # Skip 'cargo test' boilerplate at the beginning:
        lines = output.splitlines()
        start_index = 0
        for i, line in enumerate(lines):
            if f"{test_name} stdout" in line:
                start_index = i
                break

        for line in lines[start_index + 1 :]:
            if "MDTEST_TEST_FILTER" in line:
                continue
            if line.strip() == "-" * 50:
                # Skip 'cargo test' boilerplate at the end
                break

            print(line)

    def watch(self):
        def keyboard_input() -> None:
            for _ in sys.stdin:
                # This is silly, but there is no other way to inject events into
                # the main `watch` loop. We use changes to the `README.md` file
                # as a trigger to re-run all mdtests:
                MDTEST_README.touch()

        input_thread = threading.Thread(target=keyboard_input, daemon=True)
        input_thread.start()

        self._recompile_tests("Compiling tests...", message_on_success=False)
        self._run_mdtest(self.filters)
        self.console.print("[dim]Ready to watch for changes...[/dim]")

        for changes in watch(*DIRS_TO_WATCH):
            new_md_files = set()
            changed_md_files = set()
            rust_code_has_changed = False
            vendored_typeshed_has_changed = False

            for change, path_str in changes:
                path = Path(path_str)

                # See above: `README.md` changes trigger a full re-run of all tests
                if path == MDTEST_README:
                    self._run_mdtest(self.filters)
                    continue

                match path.suffix:
                    case ".rs":
                        rust_code_has_changed = True
                    case ".pyi" if path.is_relative_to(TY_VENDORED):
                        vendored_typeshed_has_changed = True
                    case ".md":
                        pass
                    case _:
                        continue

                try:
                    relative_path = Path(path).relative_to(MDTEST_DIR)
                except ValueError:
                    continue

                match change:
                    case Change.added:
                        # When saving a file, some editors (looking at you, Vim) might first
                        # save the file with a temporary name (e.g. `file.md~`) and then rename
                        # it to the final name. This creates a `deleted` and `added` change.
                        # We treat those files as `changed` here.
                        if (Change.deleted, path_str) in changes:
                            changed_md_files.add(relative_path)
                        else:
                            new_md_files.add(relative_path)
                    case Change.modified:
                        changed_md_files.add(relative_path)
                    case Change.deleted:
                        # No need to do anything when a Markdown test is deleted
                        pass
                    case _ as unreachable:
                        assert_never(unreachable)

            if rust_code_has_changed:
                if self._recompile_tests("Rust code has changed, recompiling tests..."):
                    self._run_mdtest(self.filters)
            elif vendored_typeshed_has_changed and self._recompile_tests(
                "Vendored typeshed has changed, recompiling tests..."
            ):
                self._run_mdtest(self.filters)

            for path in new_md_files | changed_md_files:
                self._run_mdtests_for_file(path)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="A runner for Markdown-based tests for ty"
    )
    parser.add_argument(
        "filters",
        nargs="*",
        help="Partial paths or mangled names, e.g., 'loops/for.md' or 'loops_for'",
    )
    parser.add_argument(
        "--enable-external",
        "-e",
        action="store_true",
        help="Enable tests with external dependencies",
    )
    parser.add_argument(
        "--no-lockfile-upgrades",
        action="store_true",
        help="By default, lockfiles will be upgraded when dependency requirements in the Markdown test change."
        + " Set this flag to never upgrade any lockfiles.",
    )

    args = parser.parse_args()

    try:
        runner = MDTestRunner(
            filters=args.filters,
            enable_external=args.enable_external,
            upgrade_lockfiles=not args.no_lockfile_upgrades,
        )
        runner.watch()
    except KeyboardInterrupt:
        print()


if __name__ == "__main__":
    main()