Source code for mathxlab.exp.logging

"""Logging utilities for experiment runs.

The logging policy is tuned for interactive experimentation:

- INFO+ from all libraries (progress, warnings)
- DEBUG only from this repository's loggers (by prefix), enabled via `verbose`

This keeps third-party debug output (e.g. Matplotlib) from overwhelming logs
while still allowing deep debugging for `mathxlab.*`.
"""

from __future__ import annotations

import contextlib
import logging
import sys
from dataclasses import dataclass
from pathlib import Path
from types import TracebackType
from typing import override

__all__ = [
    "LoggingConfig",
    "get_logger",
    "setup_logging",
]


# ------------------------------------------------------------------------------


@dataclass(frozen=True, slots=True)
class LoggingConfig:
    """
    Configuration for experiment logging.

    Args:
        package_prefix: Logger name prefix considered "our code" (default: "mathxlab").
        verbose: If True, emit DEBUG logs for package_prefix only.
        log_file: Optional file path for logs (UTF-8). The file is overwritten each run.

    Examples:
        >>> from mathxlab.exp.logging import LoggingConfig
        >>> LoggingConfig  # doctest: +SKIP
    """

    package_prefix: str = "mathxlab"
    verbose: bool = False
    log_file: Path | None = None



# ------------------------------------------------------------------------------
class _PrefixAndExactLevelFilter(logging.Filter):
    """Allow records only if the logger name matches a prefix and level matches exactly."""

    def __init__(self, *, prefix: str, levelno: int) -> None:
        super().__init__()
        self._prefix = prefix
        self._levelno = levelno

    @override
    def filter(self, record: logging.LogRecord) -> bool:
        """Return True if the record should be logged."""
        return record.name.startswith(self._prefix) and record.levelno == self._levelno


# ------------------------------------------------------------------------------


def setup_logging(*, config: LoggingConfig | None = None) -> None:
    """
    Set up logging for experiment runs.

    Design goal:
        - Show INFO+ from all libraries (progress + warnings).
        - Show DEBUG only from *our* code (by logger name prefix).

    This prevents very noisy third-party DEBUG output (e.g. matplotlib/PIL) while still
    allowing you to turn on detailed debugging for the repository code.

    Args:
        config: Logging configuration. If omitted, defaults are used.

    Examples:
        >>> from mathxlab.exp.logging import LoggingConfig, setup_logging
        >>> setup_logging(config=LoggingConfig(verbose=True))
        >>> import logging; logging.getLogger("mathxlab").debug("debug enabled")
    """
    cfg = config or LoggingConfig()

    # Ensure our console handlers can handle Unicode safely on Windows consoles.
    # PowerShell/Windows console defaults can be cp1252; that can raise
    # UnicodeEncodeError for symbols like χ, φ, π in log messages. Using
    # backslashreplace keeps output readable without forcing a code page.
    for _stream in (sys.stdout, sys.stderr):
        with contextlib.suppress(Exception):
            if hasattr(_stream, "reconfigure"):
                _stream.reconfigure(errors="backslashreplace")

    root = logging.getLogger()
    root.handlers.clear()

    # Let handlers decide what is emitted. We want to allow DEBUG to reach the DEBUG handler.
    root.setLevel(logging.DEBUG)

    fmt = logging.Formatter(
        fmt="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )

    # 1) General handler: INFO+ from everyone
    h_info = logging.StreamHandler(sys.stdout)
    h_info.setLevel(logging.INFO)
    h_info.setFormatter(fmt)
    root.addHandler(h_info)

    # 2) Optional debug handler: DEBUG only from our code
    if cfg.verbose:
        h_dbg = logging.StreamHandler(sys.stdout)
        h_dbg.setLevel(logging.DEBUG)
        h_dbg.setFormatter(fmt)
        h_dbg.addFilter(
            _PrefixAndExactLevelFilter(prefix=cfg.package_prefix, levelno=logging.DEBUG)
        )
        root.addHandler(h_dbg)

    # Optional file logging (same policy)
    if cfg.log_file is not None:
        cfg.log_file.parent.mkdir(parents=True, exist_ok=True)

        fh_info = logging.FileHandler(cfg.log_file, mode="w", encoding="utf-8")
        fh_info.setLevel(logging.INFO)
        fh_info.setFormatter(fmt)
        root.addHandler(fh_info)

        if cfg.verbose:
            fh_dbg = logging.FileHandler(cfg.log_file, mode="a", encoding="utf-8")
            fh_dbg.setLevel(logging.DEBUG)
            fh_dbg.setFormatter(fmt)
            fh_dbg.addFilter(
                _PrefixAndExactLevelFilter(prefix=cfg.package_prefix, levelno=logging.DEBUG)
            )
            root.addHandler(fh_dbg)

    # If a log file is configured, also capture unhandled exceptions into it.
    # This makes log files useful even when running without shell redirection.
    if cfg.log_file is not None:

        def _excepthook(
            exc_type: type[BaseException],
            exc: BaseException,
            tb: TracebackType | None,
        ) -> None:
            if issubclass(exc_type, KeyboardInterrupt):
                sys.__excepthook__(exc_type, exc, tb)
                return
            logging.getLogger(cfg.package_prefix).critical(
                "Unhandled exception",
                exc_info=(exc_type, exc, tb),
            )
            sys.__excepthook__(exc_type, exc, tb)

        sys.excepthook = _excepthook

    # Keep typical noisy libraries at WARNING+ (defensive). This does not hide warnings/errors.
    for noisy in ("matplotlib", "PIL"):
        logging.getLogger(noisy).setLevel(logging.WARNING)



# ------------------------------------------------------------------------------


def get_logger(name: str) -> logging.Logger:
    """
    Get a logger for a specific module.

    Args:
        name: Name of the module.

    Returns:
        A logging.Logger instance.

    Examples:
        >>> from mathxlab.exp.logging import get_logger
        >>> get_logger  # doctest: +SKIP
    """
    return logging.getLogger(name)