rich-print/src/rp/detect.py

"""File type detection for rp."""

from __future__ import annotations

import json
import shlex
from numbers import Real
from pathlib import Path

from pygments.lexers import guess_lexer
from pygments.util import ClassNotFound

EXTENSION_MAP: dict[str, str] = {
    ".md": "markdown",
    ".json": "json",
    ".py": "python",
    ".pyw": "python",
    ".js": "javascript",
    ".mjs": "javascript",
    ".ts": "typescript",
    ".tsx": "typescript",
    ".jsx": "jsx",
    ".css": "css",
    ".scss": "scss",
    ".html": "html",
    ".htm": "html",
    ".xml": "xml",
    ".yaml": "yaml",
    ".yml": "yaml",
    ".toml": "toml",
    ".sh": "bash",
    ".bash": "bash",
    ".zsh": "bash",
    ".sql": "sql",
    ".rs": "rust",
    ".go": "go",
    ".java": "java",
    ".c": "c",
    ".h": "c",
    ".cpp": "cpp",
    ".hpp": "cpp",
    ".cc": "cpp",
    ".rb": "ruby",
    ".php": "php",
    ".swift": "swift",
    ".kt": "kotlin",
    ".r": "r",
    ".lua": "lua",
    ".dockerfile": "docker",
    ".ini": "ini",
    ".cfg": "ini",
    ".diff": "diff",
    ".patch": "diff",
    ".tex": "latex",
}

_FILENAME_MAP: dict[str, str] = {
    "Dockerfile": "docker",
    "Makefile": "make",
    ".bashrc": "bash",
    ".zshrc": "bash",
    ".profile": "bash",
}

_SHEBANG_MAP: dict[str, str] = {
    "python": "python",
    "python3": "python",
    "python2": "python",
    "sh": "bash",
    "bash": "bash",
    "zsh": "bash",
}

_ENV_OPTIONS_WITH_ARG = {
    "-C",
    "-u",
    "--argv0",
    "--block-signal",
    "--chdir",
    "--default-signal",
    "--ignore-signal",
    "--unset",
}

_ENV_LONG_OPTIONS_WITH_ARG = {
    opt.lstrip("-") for opt in _ENV_OPTIONS_WITH_ARG if opt.startswith("--")
}


def _is_env_assignment(part: str) -> bool:
    """Return whether a token looks like an env-style variable assignment."""
    name, separator, _ = part.partition("=")
    return (
        bool(separator)
        and bool(name)
        and (name[0].isalpha() or name[0] == "_")
        and all(char.isalnum() or char == "_" for char in name[1:])
    )


def _detect_env_program(parts: list[str]) -> str | None:
    """Return the actual program name from an env-based shebang."""
    remaining = parts[1:]

    while remaining:
        part = remaining.pop(0)

        if part.startswith("--split-string="):
            try:
                remaining = shlex.split(part.partition("=")[2]) + remaining
            except ValueError:
                return None
            continue

        if part in {"-S", "--split-string"}:
            if not remaining:
                return None
            try:
                remaining = shlex.split(remaining.pop(0)) + remaining
            except ValueError:
                return None
            continue

        if part.startswith("-S"):
            try:
                remaining = shlex.split(part[2:]) + remaining
            except ValueError:
                return None
            continue

        if part == "--":
            if not remaining:
                return None
            return Path(remaining[0]).name

        if part in _ENV_OPTIONS_WITH_ARG:
            if not remaining:
                return None
            remaining.pop(0)
            continue

        if any(part.startswith(f"{option}=") for option in _ENV_LONG_OPTIONS_WITH_ARG):
            continue

        if part.startswith("-"):
            continue

        if _is_env_assignment(part):
            continue

        return Path(part).name

    return None


def _detect_shebang(content: str) -> str | None:
    """Return a lexer from the first shebang line when recognized."""
    lines = content.splitlines()
    first_line = lines[0] if lines else content
    if not first_line.startswith("#!"):
        return None

    command = first_line[2:].strip()
    if not command:
        return None

    try:
        parts = shlex.split(command)
    except ValueError:
        return None

    if not parts:
        return None

    program = Path(parts[0]).name
    if program == "env":
        env_program = _detect_env_program(parts)
        if env_program is None:
            return None
        program = env_program

    return _SHEBANG_MAP.get(program)


def _detect_pygments_content(content: str) -> str | None:
    """Return a lexer guessed from content when Pygments is confident."""
    try:
        lexer = guess_lexer(content)
    except ClassNotFound:
        return None

    aliases: list[str] = getattr(lexer, "aliases", [])
    alias = aliases[0] if aliases else None
    analyse_text = getattr(lexer, "analyse_text", None)
    score_value = analyse_text(content) if callable(analyse_text) else 0.0
    score = float(score_value) if isinstance(score_value, Real) else 0.0
    if alias == "text" or score < 0.5:
        return None
    return alias


def detect_type(
    path: str | None = None,
    content: str | None = None,
    explicit_type: str | None = None,
    guess_content: bool = False,
) -> str:
    """Detect the file type for syntax highlighting.

    Detection priority:
    1. Explicit type if provided
    2. Filename-based mapping (for example, Dockerfile)
    3. Extension-based mapping
    4. Shebang-based mapping for script content
    5. JSON if content starts with '{' or '[' and parses successfully
    6. Heuristic content guessing when enabled
    7. Fallback to 'text'

    Args:
        path: File path for filename and extension detection.
        content: File content for JSON detection.
        explicit_type: Override auto-detection with explicit type.
        guess_content: Enable heuristic content-based lexer guessing.

    Returns:
        Pygments lexer name for the detected file type.
    """
    if explicit_type is not None:
        return explicit_type

    if path is not None:
        p = Path(path)
        name = p.name
        if name in _FILENAME_MAP:
            return _FILENAME_MAP[name]
        suffix = p.suffix.lower()
        if suffix in EXTENSION_MAP:
            return EXTENSION_MAP[suffix]

    if content is not None:
        shebang_type = _detect_shebang(content)
        if shebang_type is not None:
            return shebang_type

        stripped = content.strip()
        if stripped and (
            (stripped[0] == "{" and stripped[-1] == "}")
            or (stripped[0] == "[" and stripped[-1] == "]")
        ):
            try:
                json.loads(stripped)
                return "json"
            except (json.JSONDecodeError, ValueError):
                pass

        if guess_content:
            pygments_type = _detect_pygments_content(content)
            if pygments_type is not None:
                return pygments_type

    return "text"