256-color Pygments highlighter with automatic light/dark background

highlight256.py

#!/usr/bin/env python3
import argparse
import sys
import shutil
from typing import Optional, Tuple, List, Any, Iterable, Type

from pygments import lex
from pygments.lexers import get_lexer_for_filename, TextLexer
from pygments.styles import get_style_by_name, get_all_styles
from pygments.token import Token
from pygments.style import Style
from pygments.lexer import Lexer


def hex_to_rgb(hex_str: str) -> Tuple[int, int, int]:
    """Convert a hex color string to an RGB tuple.

    Args:
        hex_str (str): A 6-character hexadecimal color string, with or without leading '#'.

    Returns:
        Tuple[int, int, int]: A tuple representing (R, G, B). Returns (0, 0, 0) if input is invalid.
    """
    hex_str = hex_str.lstrip("#") if hex_str else ""
    if len(hex_str) != 6:
        return (0, 0, 0)
    return tuple(int(hex_str[i : i + 2], 16) for i in (0, 2, 4))


def rgb_to_ansi256(r: int, g: int, b: int) -> int:
    """Map an RGB color to the nearest ANSI 256-color palette index.

    Args:
        r (int): Red component (0-255).
        g (int): Green component (0-255).
        b (int): Blue component (0-255).

    Returns:
        int: ANSI 256-color index.
    """
    r6 = round(r / 255 * 5)
    g6 = round(g / 255 * 5)
    b6 = round(b / 255 * 5)
    return 16 + 36 * r6 + 6 * g6 + b6


def is_light_color(hex_str: str) -> bool:
    """Determine whether a hex color is light based on relative luminance.

    Args:
        hex_str (str): A hex color string.

    Returns:
        bool: True if the color is light, False otherwise.
    """
    r, g, b = hex_to_rgb(hex_str)
    lum = 0.2126 * r + 0.7152 * g + 0.0722 * b
    return lum > 100


def get_text_color_hex(style: Type[Style]) -> Optional[str]:
    """Extract the main text color from a Pygments style.

    Args:
        style (Type[Style]): A Pygments Style class.

    Returns:
        Optional[str]: The hex color string for text, or None if not defined.
    """

    # Based on inspection of several Pygments styles:
    # - In light themes, the "color" field for Token.Name or Token.Text is often set
    #   to a dark hex value (e.g., "000000").
    # - In dark themes, the "color" field for these tokens is often omitted,
    #   resulting in None.

    name_spec = style.style_for_token(Token.Name)
    text_hex = name_spec.get("color")
    if not text_hex:
        text_spec = style.style_for_token(Token.Text)
        text_hex = text_spec.get("color")
    return text_hex


def highlight_lines(
    tokens: Iterable[Tuple[Any, str]],
    style: Type[Style],
    default_fg: int,
    background_colors: Optional[Tuple[int, int]] = None,
    fill_background: bool = False,
) -> None:
    """Highlight lines of code with optional alternating background colors.

    Args:
        tokens (Iterable[Tuple[Any, str]]): Pygments token stream.
        style (Type[Style]): Pygments style class.
        default_fg (int): Fallback ANSI color index for text.
        background_colors (Optional[Tuple[int, int]]): Pair of ANSI background colors (even, odd).
        fill_background (bool): Whether to fill each line to terminal width.
    """
    if background_colors is not None:
        bg_even, bg_odd = background_colors
    else:
        bg_even = bg_odd = None

    term_width: Optional[int] = None
    if fill_background:
        try:
            term_width = shutil.get_terminal_size().columns
        except:
            pass

    lines: List[List[Tuple[Any, str]]] = []
    current: List[Tuple[Any, str]] = []
    for ttype, val in tokens:
        parts = val.split("\n")
        for i, part in enumerate(parts):
            if part:
                current.append((ttype, part))
            if i < len(parts) - 1:
                lines.append(current)
                current = []
    if current:
        lines.append(current)

    for lineno, line_tokens in enumerate(lines, start=1):
        bg = bg_even if lineno % 2 == 0 else bg_odd
        if bg is not None:
            line_output = f"\033[48;5;{bg}m"
        else:
            line_output = ""
        text_len = 0
        for ttype, part in line_tokens:
            spec = style.style_for_token(ttype)
            hexcol = spec.get("color")
            if hexcol:
                r, g, b = hex_to_rgb(hexcol)
                fg_idx = rgb_to_ansi256(r, g, b)
            else:
                fg_idx = default_fg
            line_output += f"\033[38;5;{fg_idx}m{part}\033[39m"
            text_len += len(part)
        if term_width is not None:
            pad_len = term_width - text_len
            if pad_len > 0:
                line_output += " " * pad_len
        if bg is not None:
            line_output += "\033[0m\n"
        sys.stdout.write(line_output)


def main() -> None:
    available_styles: List[str] = list(get_all_styles())

    # Check for --list-styles manually before parsing
    if "--list-styles" in sys.argv:
        print("Available styles:")
        for style_name in sorted(available_styles):
            print(f"  {style_name}")
        sys.exit(0)

    parser = argparse.ArgumentParser(description="256-color Pygments highlighter with automatic light/dark background")
    parser.add_argument("filename", help="Path to source file to highlight")
    parser.add_argument("--style", default="default", help="Pygments style to use (see --list-styles)")
    parser.add_argument("--fill-background", action="store_true", help="Fill the line background to the right end")
    parser.add_argument("--list-styles", action="store_true", help="List available Pygments styles and exit")
    args = parser.parse_args()

    assert not args.list_styles

    if args.style not in available_styles:
        parser.error(f"Unknown style '{args.style}'. Use --list-styles to see available styles.")

    try:
        with open(args.filename, encoding="utf-8") as f:
            code: str = f.read()
    except Exception as e:
        print(f"Error: Fail to read '{args.filename}': {e}", file=sys.stderr)
        sys.exit(1)

    try:
        lexer: Lexer = get_lexer_for_filename(args.filename, code)
    except Exception:
        print(f"Warning: No lexer for '{args.filename}', using plain text.", file=sys.stderr)
        lexer = TextLexer()

    style_cls: Type[Style] = get_style_by_name(args.style)
    text_hex: str = get_text_color_hex(style_cls) or "ffffff"

    r0, g0, b0 = hex_to_rgb(text_hex)
    default_fg: int = rgb_to_ansi256(r0, g0, b0)

    if is_light_color(text_hex):
        bg_odd, bg_even = 234, 235
    else:
        bg_odd, bg_even = 254, 253

    bg_colors: Tuple[int, int] = (bg_even, bg_odd)
    tokens = list(lex(code, lexer))
    highlight_lines(tokens, style_cls, default_fg, background_colors=bg_colors, fill_background=args.fill_background)


if __name__ == "__main__":
    main()