A loguru logger that uses rich to print to the console.

rich_logger.py

"""Create a rich.console.Console sink for a loguru.Logger.

This module provides a custom sink for the loguru logger that uses the rich library to format and display log messages in a visually appealing way. It also includes functions to set up logging, manage run counts, and handle log file creation and cleanup.
It is designed to be used in Python projects that require structured and colorful logging output, especially for long-running processes or applications where tracking progress and errors is important.

The module includes:
- A custom RichSink class that formats log messages with colors and styles.
- Functions to create and configure a rich console and progress bar.
- Functions to find the current working directory and manage log files.
- A function to increment and write the run count to a file.
- A function to handle cleanup and log file management on exit.
"""

from __future__ import annotations

import atexit
import re
import sys
from collections import deque
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, cast

import loguru
from rich.console import Console
from rich.panel import Panel
from rich.progress import (
    BarColumn,
    MofNCompleteColumn,
    Progress,
    SpinnerColumn,
    TextColumn,
    TimeElapsedColumn,
    TimeRemainingColumn,
)
from rich.style import Style
from rich.text import Text
from rich.traceback import install as tr_install
from rich_gradient import Gradient
from rich_gradient.color import Color
from rich_gradient.rule import GradientRule

__all__ = [
    "get_console",
    "get_logger",
    "get_progress",
    "get_progress_console",
    "find_cwd",
    "CWD",
    "LOGS_DIR",
    "RUN_FILE",
    "FORMAT",
    "trace_sink",
    "RichSink",
    "on_exit",
    "setup",
    "read_run_from_file",
    "write_run_to_file",
    "increment_run_and_write_to_file",
]

__version__ = "0.2.1"
__author__ = "Max Ludden"

VERBOSE: bool = False
TRACK_RUN: bool = True
HandlerConfig = Dict[str, Any]

# Log level names for validation
_LEVEL_NAMES = [
    "TRACE",
    "DEBUG",
    "INFO",
    "SUCCESS",
    "WARNING",
    "ERROR",
    "CRITICAL",
]


def get_progress(console: Optional[Console] = None) -> Progress:
    """
    Initialize and return a Rich progress bar.

    Args:
        console (Optional[Console]): An optional existing Rich console.
    Returns:
        Progress: A configured Rich progress bar.
    """
    if console is None:
        console = Console()
    progress = Progress(
        SpinnerColumn(spinner_name="point"),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(bar_width=None),
        TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
        TimeElapsedColumn(),
        TimeRemainingColumn(),
        MofNCompleteColumn(),
        console=console,
        expand=True,
        refresh_per_second=30,
        transient=True,
    )
    return progress


def get_console(
    console: Optional[Console] = None, record: bool = False, show_locals: bool = False
) -> Console:
    """
    Initialize and return a Rich console.

    Args:
        record (bool): Whether to record console output.
        show_locals (bool): Whether to show local variables in tracebacks.
        console (Optional[Console]): An optional existing Rich console (unused).

    Returns:
        Console: A configured Rich console.
    """
    # The `console` argument is unused and always overridden.
    global _console

    if console is not None:
        _console = console
    else:
        _console = Console(record=record)
    tr_install(console=_console, show_locals=show_locals)
    return _console


def get_progress_console(
    console: Optional[Console] = None,
    record: bool = False,
    show_locals: bool = False,
) -> Tuple[Console, Progress]:
    """
    Convenience function to get both a console and a progress bar.

    Args:
        console (Optional[Console]): An optional existing Rich console.
        record (bool): Whether to record console output.
        show_locals (bool): Whether to show local variables in tracebacks.
    Returns:
        Tuple[Console, Progress]: A tuple containing a configured Rich console and a progress bar.
    """
    if console is None:
        console = get_console(record=record, show_locals=show_locals)
    progress = get_progress(console=console)
    console = progress.console
    return console, progress


def find_cwd(
    start_dir: Path = Path(__file__).parent.parent, verbose: bool = False
) -> Path:
    """
    Find the current working directory by walking upward until a 'pyproject.toml' is found.

    Args:
        start_dir (Path): The starting directory.
        verbose (bool): If True, prints the found directory in a styled panel.

    Returns:
        Path: The current working directory.
    """
    cwd: Path = start_dir
    while not (cwd / "pyproject.toml").exists():
        cwd = cwd.parent
        if cwd == Path.home():
            break
    if verbose:
        console = get_console()
        console.line(2)
        panel_title = Gradient(
            "Current Working Directory",
            colors=[
                Color("#ff005f"),
                Color("#ff00af"),
                Color("#ff00ff"),
            ],
            style="bold",
        ).as_text()
        console.print(
            Panel(
                f"[i #5f00ff]{cwd.resolve()}",
                title=panel_title,
            )
        )
        console.line(2)
    return cwd


# Constants and paths
CWD: Path = find_cwd()
LOGS_DIR: Path = CWD / "logs"
RUN_FILE: Path = LOGS_DIR / "run.txt"
FORMAT: str = (
    "{time:hh:mm:ss.SSS} | {file.name: ^12} | Line {line} | {level} ➤ {message}"
)


def trace_sink() -> Dict[str, Any]:
    """
    Return the configuration for the trace sink.
    Returns:
        Dict[str, Any]: The trace sink configuration.
    """
    return {
        "sink": str((LOGS_DIR / "trace.log").resolve()),
        "format": FORMAT,
        "level": "TRACE",
        "backtrace": True,
        "diagnose": True,
        "colorize": False,
        "mode": "w",
    }


def setup() -> Optional[int]:
    """
    Setup the logger by creating necessary directories and files.

    Returns:
        Optional[int]: The run count (read from the run file), or None if run tracking is disabled.
    """
    console = get_console()
    if not LOGS_DIR.exists():
        LOGS_DIR.mkdir(parents=True)
        console.print(f"Created Logs Directory: {LOGS_DIR}")
    if not TRACK_RUN:
        console.print("[i #af99ff]Setup logger. Disabling run tracking.[/]")
        return None
    if not RUN_FILE.exists():
        with open(RUN_FILE, "w", encoding="utf-8") as f:
            f.write("0")
            console.print("Created Run File, set to 0")
    with open(RUN_FILE, "r", encoding="utf-8") as f:
        run = int(f.read())
    return run


def read_run_from_file() -> Optional[int]:
    """
    Read the run count from the run file.

    Returns:
        Optional[int]: The run count, or None if tracking is disabled.
    """
    if not TRACK_RUN:
        return None
    console = get_console()
    if not RUN_FILE.exists():
        console.print("[b #ff0000]Run File Not Found[/][i #ff9900] – Creating...[/]")
        setup()
    with open(RUN_FILE, "r", encoding="utf-8") as f:
        run = int(f.read())
    return run


def get_default_sinks(
    console: Optional[Console],
    run: Optional[int],
    level: int,
    padding: Tuple[int, int] = (0, 1),
    expand: bool = False
) -> List[HandlerConfig]:
    """
    Return the default sinks for the logger.
    Args:
        console (Optional[Console]): The Rich console for the RichSink.
        run (Optional[int]): The run count.
        level (int): The log level as an integer.
    Returns:
        List[Dict[str, Any]]: A list of sink configuration dictionaries.
    """
    return [
        {
            "sink": RichSink(
                console=console,
                run=run,
                record=False,
                padding=padding,
                expand=expand
            ),
            "format": "{message}",
            "level": level,
            "backtrace": True,
            "diagnose": True,
            "colorize": False,
        },
        {
            "sink": str(LOGS_DIR / "trace.log"),
            "format": FORMAT,
            "level": "TRACE",
            "backtrace": True,
            "diagnose": True,
            "colorize": False,
            "mode": "a",  # Use append mode instead of write mode
            "retention": "30 minutes",
        },
    ]


def _validate_level(level: str | int) -> int:
    """
    Validate the log level and convert it to an integer.

    Args:
        level (str|int): The logging level. Can be a string (e.g., "DEBUG", "INFO", etc.) or an integer (0-50).

    Returns:
        Optional[int]: The validated log level as an integer, or None if invalid.

    Raises:
        TypeError: If the log level is not a string or an integer.
        ValueError: If the log level is not valid.
    """
    if isinstance(level, int):
        if not (0 <= level <= 50):
            raise ValueError(
                f"Log level integer must be between 0 and 50, got {level}."
            )
        return level
    if not isinstance(level, str):
        raise TypeError(f"Log level must be a string or an integer, got {type(level)}.")
    _level = level.upper()
    if _level not in _LEVEL_NAMES:
        raise ValueError(
            f"Invalid log level: {level!r}. Must be one of: {', '.join(_LEVEL_NAMES)}."
        )
    match _level:
        case "TRACE":
            return 5
        case "DEBUG":
            return 10
        case "INFO":
            return 20
        case "SUCCESS":
            return 25
        case "WARNING":
            return 30
        case "ERROR":
            return 40
        case "CRITICAL":
            return 50
        case _:
            raise ValueError(
                f"Unable to parse log level: {level!r}. Must be one of: {', '.join(_LEVEL_NAMES)}."
            )


def get_logger(
    console: Optional[Console] = None,
    level: str | int = "DEBUG",
    verbose: bool = False,
    track_run: bool = True,
    additional_sinks: Optional[List[Dict[str, Any]]] = None,
    padding: Tuple[int, int] = (1, 2),
    expand: bool = True
) -> loguru.Logger:
    """
    Initialize and return a Loguru logger.

    Args:
        console (Optional[Console]): An optional existing Rich console. If None, a new one is created.
        level (str|int): The logging level (e.g., "DEBUG", "INFO", etc.) or an integer (0-50).
        verbose (bool): If True, enables verbose logging.
        track_run (bool): If True, enables run tracking.
        additional_sinks (Optional[List[Dict[str, Any]]]): Extra sinks to add to the logger.

    Returns:
        Logger: A configured Loguru logger.
    """
    global VERBOSE
    VERBOSE = verbose  # Store the value for other functions to access
    global TRACK_RUN
    TRACK_RUN = track_run  # Store the value for other functions to access
    if console is None:
        console = get_console()
    run = read_run_from_file() if track_run else None

    # Validate the log level
    _level = _validate_level(level)
    sinks = get_default_sinks(console=console, run=run, level=_level, padding=padding, expand=expand)
    if additional_sinks:
        sinks.extend(additional_sinks)

    log = loguru.logger.bind(sink="rich")
    log.remove()
    log.configure(
        handlers=cast(Any, sinks),
        extra={"run": run, "rich": "", "verbose": verbose, "padding": ()},
    )
    return log


def write_run_to_file(run: int, verbose: bool = False) -> None:
    """
    Write the run count to the run file.

    Args:
        run (int): The run count to write.
        verbose (bool): If True, logs a trace message.
    """
    if verbose:
        log = get_logger()
        log.trace("Writing run count...")
    with open(RUN_FILE, "w", encoding="utf-8") as f:
        f.write(str(run))


def increment_run_and_write_to_file() -> Optional[int]:
    """
    Increment the run count, write it to the file, and return the new count.

    Returns:
        Optional[int]: The incremented run count, or None if tracking is disabled.
    """
    if not TRACK_RUN:
        return None
    log = get_logger()
    log.trace("Incrementing run count...")
    run = read_run_from_file()
    assert run is not None, "Run count not found in file."
    run += 1
    write_run_to_file(run)
    return run


class RichSink:
    """
    A custom Loguru sink that uses Rich to print styled log messages.

    Attributes:
        LEVEL_STYLES (Dict[str, Style]): Styles for each log level.
        GRADIENTS (Dict[str, List[Color]]): Gradients for log level titles.
        MSG_COLORS (Dict[str, List[Color]]): Gradients for log message text.
        run (Optional[int]): The current run number.
        console (Console): The Rich console used for output.
    """

    LEVEL_STYLES: Dict[str, Style] = {
        "TRACE": Style(italic=True),
        "DEBUG": Style(color="#aaaaaa"),
        "INFO": Style(color="#00afff"),
        "SUCCESS": Style(bold=True, color="#00ff00"),
        "WARNING": Style(italic=True, color="#ffaf00"),
        "ERROR": Style(bold=True, color="#ff5000"),
        "CRITICAL": Style(bold=True, color="#ff0000"),
    }
    """Styles for each log level."""
    GRADIENTS: Dict[str, list[Color]] = {
        "TRACE": [Color("#888888"), Color("#aaaaaa"), Color("#cccccc")],
        "DEBUG": [Color("#0F8C8C"), Color("#19cfcf"), Color("#00ffff")],
        "INFO": [Color("#1b83d3"), Color("#00afff"), Color("#54d1ff")],
        "SUCCESS": [Color("#00ff90"), Color("#00ff00"), Color("#afff00")],
        "WARNING": [Color("#ffaa00"), Color("#ffcc00"), Color("#ffff00")],
        "ERROR": [Color("#ff7700"), Color("#ff5500"), Color("#ff3300")],
        "CRITICAL": [Color("#ff0000"), Color("#ff005f"), Color("#ff009f")],
    }
    """Gradients for log level titles."""
    MSG_COLORS: Dict[str, list[Color]] = {
        "TRACE": [Color("#eeeeee"), Color("#dddddd"), Color("#bbbbbb")],
        "INFO": [Color("#a4e7ff"), Color("#72d3ff"), Color("#52daff")],
        "SUCCESS": [Color("#d3ffd3"), Color("#a9ffa9"), Color("#64ff64")],
        "WARNING": [Color("#ffeb9b"), Color("#ffe26e"), Color("#ffc041")],
        "ERROR": [Color("#ffc59c"), Color("#ffaa6e"), Color("#FF4E3A")],
        "CRITICAL": [Color("#ffaaaa"), Color("#FF6FA4"), Color("#FF49C2")],
    }
    """Gradients for log message text."""

    def __init__(
        self,
        console: Optional[Console] = None,
        track_run: bool = TRACK_RUN,
        run: Optional[int] = None,
        padding: Tuple[int, int] = (1, 2),
        expand: bool = False,
        record: bool = False
    ) -> None:
        """
        Args:
            console (Optional[Console]): An optional Rich console.
            track_run (bool): Whether to track runs.
            run (Optional[int]): The current run number. If None, it is read from the run file.
        """
        if track_run:
            if run is None:
                try:
                    run = read_run_from_file()
                except FileNotFoundError:
                    run = setup()
            self.run = run
        else:
            self.run = None
        self.console: Console = console or get_console()
        if record:
            self.console.record = True
            self.record = True
        else:
            self.console.record = False
            self.record = False
        self.padding: Tuple[int, int] = padding or (1, 2)
        self.expand: bool = expand

    def __call__(self, message: loguru.Message) -> None:
        """
        Print a loguru.Message to the Rich console as a styled panel.

        Args:
            message (Message): The loguru message to print.
        """
        record = message.record
        panel = self._build_panel(record, self.run)
        self.console.print(panel)

    def _build_panel(self, record: loguru.Record, run: Optional[int] = None) -> Panel:
        """
        Helper method to build a Rich Panel for a log record.

        Args:
            record (Record): The log record.
            run (Optional[int]): The current run count.

        Returns:
            Panel: A Rich Panel containing the formatted log message.
        """
        level_name = record["level"].name
        colors = self.GRADIENTS.get(level_name, [])
        style = self.LEVEL_STYLES.get(level_name, Style())
        msg_style = self.MSG_COLORS.get(
            level_name,
            [
                Color("#eeeeee"),
                Color("#aaaaaa"),
                Color("#888888"),
            ],
        )
        # Title with gradient and highlighted separators.
        title: Text = Gradient(
            f" {level_name} | {record['file'].name} | Line {record['line']} ",
            colors=colors,
        ).as_text()
        title.stylize(Style(reverse=True))
        # Subtitle with run count and formatted time.
        subtitle_run_elements: list[Text] = [
            Text(f"Run {run}"),
            Text(" | "),
        ]
        subtitle_elements: list[Text] = [
            Text(record["time"].strftime("%h:%M:%S.%f")[:-3]),
            Text(record["time"].strftime(" %p")),
        ]
        if run is not None:
            subtitle_elements = subtitle_run_elements + subtitle_elements
        subtitle: Text = Text.assemble(*subtitle_elements)
        subtitle.highlight_words(":", style="dim #aaaaaa")
        # Message text with gradient.
        message_text: Text = Gradient(record["message"], colors=msg_style)
        return Panel(
            message_text,
            title=title,
            title_align="left",
            subtitle=subtitle,
            subtitle_align="right",
            border_style=style + Style(bold=True),
            padding=self.padding,
            expand=self.expand
        )


_RUN_HEADER_PATTERN = re.compile(r"Run (\d+) Completed")


def on_exit() -> None:
    """
    At exit, increment the run count, add a header to the run’s log,
    and trim the trace log to the last three runs.
    """
    log = get_logger()
    run = increment_run_and_write_to_file()
    if VERBOSE:
        log.info(f"Run {run} Completed")
    segments: deque[str] = deque(maxlen=3)
    current_segment: list[str] = []
    trace_log_path = LOGS_DIR / "trace.log"
    # Process the log file line by line to build segments.
    with open(trace_log_path, "r", encoding="utf-8") as f:
        for line in f:
            current_segment.append(line)
            if _RUN_HEADER_PATTERN.search(line):
                segments.append("".join(current_segment))
                current_segment = []
    # If there's remaining content, add it to the last segment.
    if current_segment:
        if segments:
            segments[-1] += "".join(current_segment)
        else:
            segments.append("".join(current_segment))
    # For each segment, insert a header if not already present.
    updated_segments = []
    for segment in segments:
        # Check if the first non-empty line already is a header.
        stripped = segment.lstrip()
        if not stripped.startswith("===="):
            m = _RUN_HEADER_PATTERN.search(segment)
            run_number = m.group(1) if m else "Unknown"
            header = f"\n===== Run {run_number} Log =====\n"
            segment = header + segment
        updated_segments.append(segment)
    trimmed_log = "\n".join(updated_segments)
    if VERBOSE:
        log.debug(f"Trimmed trace log to the last {len(updated_segments)} run(s).")
    # Overwrite the trace log with the trimmed (and header-enhanced) content.
    trace_log_path.write_text(trimmed_log, encoding="utf-8")


atexit.register(on_exit)

if __name__ == "__main__":
    console = get_console()
    log: loguru.Logger = get_logger(console=console, level="TRACE")
    console.clear()
    console.line(2)
    console.print(
        GradientRule(
            "Loguru Logger using rich.console.Console",
        )
    )
    console.line(2)


    log.trace(
        "This is a loguru.Message logged to a rich.console.Console at Level.TRACE"
    )
    log.debug(
        "This is a loguru.Message logged to a rich.console.Console at Level.DEBUG"
    )
    log.info("This is a loguru.Message logged to a rich.console.Console at Level.INFO")
    log.success(
        "This is a loguru.Message logged to a rich.console.Console at Level.SUCCESS"
    )
    log.warning(
        "This is a loguru.Message logged to a rich.console.Console at Level.WARNING"
    )
    log.error(
        "This is a loguru.Message logged to a rich.console.Console at Level.ERROR"
    )
    log.critical(
        "This is a loguru.Message logged ato rich.console.Console at Level.CRITICAL"
    )

    sys.exit(0)