Source code for ytjobs.logging.silencer

"""Redirect or suppress process stdout/stderr for mapper-style jobs."""

import os
import sys
from collections.abc import Callable, Generator
from contextlib import contextmanager
from functools import wraps
from pathlib import Path
from typing import Any, Literal, TypeVar, cast

_F = TypeVar("_F", bound=Callable[..., Any])




def manage_output(
    mode: Literal["suppress", "redirect"] = "redirect",
) -> Callable[[_F], _F]:
    """Return a decorator that suppresses output or redirects stdout to stderr.

    Args:
        mode: ``"suppress"`` for full silence, ``"redirect"`` to send stdout to stderr
            (keeps job stdout clean for JSON lines).

    Returns:
        A decorator, e.g. ``@manage_output(mode="redirect")`` above a function.

    """

    def decorator(func: _F) -> _F:
        @wraps(func)
        def wrapper(*args: object, **kwargs: object) -> object:
            if mode == "suppress":
                with suppress_all_output():
                    return func(*args, **kwargs)
            elif mode == "redirect":
                with redirect_stdout_to_stderr():
                    return func(*args, **kwargs)
            else:
                msg = f"Invalid mode: {mode}. Must be 'suppress' or 'redirect'"
                raise ValueError(msg)

        return cast("_F", wrapper)

    return decorator





@contextmanager
def redirect_stdout_to_stderr() -> Generator[None, None, None]:
    """Context manager that redirects stdout to stderr.

    This is useful for YTsaurus mappers where you need clean JSON on stdout,
    but processing functions might print debug messages.

    Usage:
        with redirect_stdout_to_stderr():
            print("This goes to stderr")  # Won't corrupt stdout
            some_function_that_prints()
    """
    original_stdout = sys.stdout
    try:
        sys.stdout = sys.stderr
        yield
    finally:
        sys.stdout = original_stdout





@contextmanager
def suppress_all_output() -> Generator[None, None, None]:
    """Context manager that suppresses ALL output: stdout, stderr, and warnings.

    This is useful when you need complete silence from noisy libraries
    like OpenCV, Ultralytics YOLO, TensorFlow, Ceres Solver, etc.

    Usage:
        with suppress_all_output():
            # All prints, warnings, and library output are suppressed
            model = YOLO('yolov8n.pt')  # No output
            results = model(image)       # No progress bars
            cv2.imread(path)             # No warnings
    """
    # Save original Python streams
    original_stdout = sys.stdout
    original_stderr = sys.stderr

    # Save original OS-level file descriptors
    original_stdout_fd = os.dup(1)  # stdout FD
    original_stderr_fd = os.dup(2)  # stderr FD

    with Path(os.devnull).open("w") as devnull:
        devnull_fd = devnull.fileno()
        try:
            # Redirect Python streams to devnull
            sys.stdout = devnull
            sys.stderr = devnull

            # Redirect OS-level file descriptors (this catches C library output like Ceres)
            os.dup2(devnull_fd, 1)  # Redirect stdout FD
            os.dup2(devnull_fd, 2)  # Redirect stderr FD

            yield

        finally:
            # Restore OS-level file descriptors first
            os.dup2(original_stdout_fd, 1)
            os.dup2(original_stderr_fd, 2)

            # Close duplicate FDs
            os.close(original_stdout_fd)
            os.close(original_stderr_fd)

            # Restore Python streams
            sys.stdout = original_stdout
            sys.stderr = original_stderr