From 07d1be8c817a4926e5497096daaa4121b96b0a1e Mon Sep 17 00:00:00 2001 From: PakitoSec Date: Sat, 4 Apr 2026 18:37:44 +0200 Subject: [PATCH] feat: register shutdown_logger with atexit for automatic log flushing on process exit --- README.md | 14 +++++++------- examples/async_context.py | 6 +----- examples/base.py | 4 +--- examples/mp_adv_data_processing.py | 3 --- examples/mp_example.py | 3 --- examples/serialize.py | 4 +--- src/logurich/core.py | 11 +++++++++++ src/logurich/struct.py | 2 ++ tests/test_core.py | 15 +++++++++++++++ tests/test_opt_click.py | 19 +++++++++++++++++++ 10 files changed, 57 insertions(+), 24 deletions(-) create mode 100644 tests/test_opt_click.py diff --git a/README.md b/README.md index 7a32133..835f4be 100644 --- a/README.md +++ b/README.md @@ -18,7 +18,7 @@ pip install logurich[click] ```python from rich.panel import Panel -from logurich import init_logger, logger, shutdown_logger +from logurich import init_logger, logger init_logger("INFO", enqueue=False) @@ -46,11 +46,12 @@ logger.info( }, ) -shutdown_logger() ``` `logger.ctx(...)` is shorthand for the existing module-level `ctx(...)` helper. `logger.contextualize(...)` is a convenience alias for `global_context_configure(...)`. The module-level helpers remain supported if you prefer `global_context_configure(...)` or `extra={"context": {"key": ctx(...)}}`. +For short-lived scripts and CLIs, `init_logger()` automatically registers an `atexit` hook, so you do not need to call `shutdown_logger()` just to flush logs at process exit. + ## Using Logurich in Reusable Libraries If you are writing a Python library that will be imported by another program, the library should not call `init_logger()` on its own. Let the main application own logging configuration, handler setup, and shutdown. @@ -88,14 +89,12 @@ Then configure Logurich once in the main program: ```python # main.py -from logurich import init_logger, shutdown_logger +from logurich import init_logger from mylib.service import run_job init_logger("INFO", enqueue=False) run_job("job-42") - -shutdown_logger() ``` Guidelines for libraries: @@ -114,7 +113,7 @@ When `enqueue=True`, Logurich is process-safe only if worker processes send reco import logging import multiprocessing as mp -from logurich import configure_child_logging, get_log_queue, init_logger, shutdown_logger +from logurich import configure_child_logging, get_log_queue, init_logger def worker(log_queue: mp.Queue, worker_id: int) -> None: @@ -137,11 +136,12 @@ def main() -> None: for process in processes: process.join() - shutdown_logger() ``` Only the process that calls `init_logger(..., enqueue=True)` owns the console and file handlers. Child processes must call `configure_child_logging(queue)` before logging. +Call `shutdown_logger()` explicitly only when you need deterministic teardown before process exit, such as in tests or when reconfiguring logging multiple times in the same interpreter. + ## Click CLI helper Install the optional Click extra to automatically expose logger configuration flags inside your commands: diff --git a/examples/async_context.py b/examples/async_context.py index cea328d..201176c 100644 --- a/examples/async_context.py +++ b/examples/async_context.py @@ -9,7 +9,6 @@ get_log_queue, global_context_configure, init_logger, - shutdown_logger, ) request_log = logging.getLogger("example.request") @@ -90,7 +89,4 @@ async def main(log_queue: mp.Queue) -> None: if __name__ == "__main__": init_logger("INFO", enqueue=True) log_queue = get_log_queue() - try: - asyncio.run(main(log_queue)) - finally: - shutdown_logger() + asyncio.run(main(log_queue)) diff --git a/examples/base.py b/examples/base.py index 2bf95c3..7f70b30 100644 --- a/examples/base.py +++ b/examples/base.py @@ -1,7 +1,7 @@ from rich.panel import Panel from rich.table import Table -from logurich import ctx, init_logger, logger, shutdown_logger +from logurich import ctx, init_logger, logger def create_rich_table() -> Table: @@ -53,5 +53,3 @@ def create_rich_table() -> Table: request_id=ctx("req-99", style="cyan", show_key=True), ) req_logger.info("Handling request") - - shutdown_logger() diff --git a/examples/mp_adv_data_processing.py b/examples/mp_adv_data_processing.py index b7ae496..209de7f 100644 --- a/examples/mp_adv_data_processing.py +++ b/examples/mp_adv_data_processing.py @@ -16,7 +16,6 @@ global_context_configure, global_context_set, init_logger, - shutdown_logger, ) @@ -164,8 +163,6 @@ def main(): extra={"renderables": (results_table, sample_results)}, ) - shutdown_logger() - if __name__ == "__main__": mp.set_start_method("spawn", force=True) diff --git a/examples/mp_example.py b/examples/mp_example.py index 72a080a..1e5d78d 100644 --- a/examples/mp_example.py +++ b/examples/mp_example.py @@ -12,7 +12,6 @@ get_log_queue, global_context_configure, init_logger, - shutdown_logger, ) @@ -109,8 +108,6 @@ def main() -> None: process.pid, ) - shutdown_logger() - if __name__ == "__main__": mp.set_start_method("spawn", force=True) diff --git a/examples/serialize.py b/examples/serialize.py index 68c4025..7ae67d1 100644 --- a/examples/serialize.py +++ b/examples/serialize.py @@ -3,7 +3,7 @@ from rich.panel import Panel from rich.table import Table -from logurich import ctx, global_context_configure, init_logger, logger, shutdown_logger +from logurich import ctx, global_context_configure, init_logger, logger def build_table() -> Table: @@ -52,5 +52,3 @@ def build_table() -> Table: raise RuntimeError("serialize example failure") except RuntimeError: logger.exception("Exception payload") - - shutdown_logger() diff --git a/src/logurich/core.py b/src/logurich/core.py index a79ebb8..990382b 100644 --- a/src/logurich/core.py +++ b/src/logurich/core.py @@ -2,6 +2,7 @@ from __future__ import annotations +import atexit import contextlib import contextvars import copy @@ -668,6 +669,15 @@ def shutdown_logger() -> None: _context_state.set({}) +def _ensure_shutdown_atexit_registered() -> None: + """Register ``shutdown_logger`` once so handlers are flushed on process exit.""" + + if logger_state.get("atexit_registered"): + return + atexit.register(shutdown_logger) + logger_state["atexit_registered"] = True + + def get_log_queue() -> mp.Queue: """Return the active multiprocessing queue used for logging.""" @@ -721,6 +731,7 @@ def init_logger( ) -> Optional[str]: """Initialize stdlib logging with optional Rich rendering and queue support.""" + _ensure_shutdown_atexit_registered() shutdown_logger() env_rich_handler = parse_bool_env("LOGURICH_RICH") diff --git a/src/logurich/struct.py b/src/logurich/struct.py index eefbf94..140f55a 100644 --- a/src/logurich/struct.py +++ b/src/logurich/struct.py @@ -9,4 +9,6 @@ "queue": None, "listener": None, "final_handlers": (), + "env_extra": {}, + "atexit_registered": False, } diff --git a/tests/test_core.py b/tests/test_core.py index fdfa7cb..c0cbda6 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -17,6 +17,7 @@ logger as exported_logger, ) from logurich.console import rich_configure_console +from logurich.struct import logger_state @pytest.mark.parametrize( @@ -33,6 +34,20 @@ def test_level_info(logger, buffer): assert "Debug, world!" not in output +def test_init_logger_registers_atexit_shutdown_once(monkeypatch): + registered: list[object] = [] + + monkeypatch.setitem(logger_state, "atexit_registered", False) + monkeypatch.setattr("logurich.core.atexit.register", registered.append) + + init_logger("INFO", enqueue=False) + shutdown_logger() + init_logger("INFO", enqueue=False) + shutdown_logger() + + assert registered == [shutdown_logger] + + @pytest.mark.parametrize( "logger", [{"level": "DEBUG", "enqueue": False}, {"level": "DEBUG", "enqueue": True}], diff --git a/tests/test_opt_click.py b/tests/test_opt_click.py new file mode 100644 index 0000000..81eb122 --- /dev/null +++ b/tests/test_opt_click.py @@ -0,0 +1,19 @@ +import pytest + +click = pytest.importorskip("click") + +from logurich import shutdown_logger # noqa: E402 +from logurich.opt_click import click_logger_init # noqa: E402 +from logurich.struct import logger_state # noqa: E402 + + +def test_click_logger_init_registers_atexit_shutdown(monkeypatch): + registered: list[object] = [] + + monkeypatch.setitem(logger_state, "atexit_registered", False) + monkeypatch.setattr("logurich.core.atexit.register", registered.append) + + click_logger_init("INFO", 0, None, (), False) + shutdown_logger() + + assert registered == [shutdown_logger]