diff --git a/src/koyo/context.py b/src/koyo/context.py index 01effa2..b75a331 100644 --- a/src/koyo/context.py +++ b/src/koyo/context.py @@ -1,8 +1,24 @@ """Various context managers.""" import gc +import os +import sys from contextlib import contextmanager +@contextmanager +def nullout(): + """Context manager to suppress stdout and stderr.""" + save_stdout = sys.stdout + save_stderr = sys.stderr + sys.stdout = open(os.devnull, "w") + sys.stderr = open(os.devnull, "w") + try: + yield + finally: + sys.stdout = save_stdout + sys.stderr = save_stderr + + @contextmanager def no_gc(): """Context manager to disable garbage collection.""" diff --git a/src/koyo/logging.py b/src/koyo/logging.py index 531e845..1763af2 100644 --- a/src/koyo/logging.py +++ b/src/koyo/logging.py @@ -76,6 +76,15 @@ def get_loguru_config( return level.upper(), fmt, colorize, enqueue +def get_stderr() -> ty.TextIO: + """Get stderr.""" + if sys.stderr is None: + sys.stderr = sys.stdout + if sys.stderr is None: + sys.stderr = open(os.devnull, "w") + return sys.stderr + + def set_loguru_log( sink=sys.stderr, level: str | int = 20, @@ -91,6 +100,12 @@ def set_loguru_log( """Set loguru formatting.""" if logger is None: from loguru import logger + if logger is None: + raise ValueError("Logger is None - cannot set loguru log.") + if sink is None: + sink = get_stderr() + if sink is None: + raise ValueError("Sink is None - cannot set loguru log.") # automatically get format fmt = fmt if fmt is not None else (LOG_FMT if no_color else COLOR_LOG_FMT)