diff --git a/src/logurich/core.py b/src/logurich/core.py index 990382b..8fec406 100644 --- a/src/logurich/core.py +++ b/src/logurich/core.py @@ -10,6 +10,7 @@ import logging.handlers import multiprocessing as mp import os +import threading import traceback from collections.abc import Mapping from dataclasses import dataclass @@ -678,6 +679,20 @@ def _ensure_shutdown_atexit_registered() -> None: logger_state["atexit_registered"] = True +def _ensure_shutdown_threading_atexit_registered() -> None: + """Register ``shutdown_logger`` before interpreter thread shutdown when possible.""" + + if logger_state.get("threading_atexit_registered"): + return + + register = getattr(threading, "_register_atexit", None) + if register is None: + return + + register(shutdown_logger) + logger_state["threading_atexit_registered"] = True + + def get_log_queue() -> mp.Queue: """Return the active multiprocessing queue used for logging.""" @@ -731,6 +746,7 @@ def init_logger( ) -> Optional[str]: """Initialize stdlib logging with optional Rich rendering and queue support.""" + _ensure_shutdown_threading_atexit_registered() _ensure_shutdown_atexit_registered() shutdown_logger() diff --git a/src/logurich/opt_click.py b/src/logurich/opt_click.py index fff6155..aa070ce 100644 --- a/src/logurich/opt_click.py +++ b/src/logurich/opt_click.py @@ -12,7 +12,7 @@ import click -from . import LOG_LEVEL_CHOICES, init_logger, logger +from . import LOG_LEVEL_CHOICES, init_logger, logger, shutdown_logger LOGGER_PARAM_NAMES = ( "logger_level", @@ -21,6 +21,7 @@ "logger_level_by_module", "logger_rich", ) +_CLICK_SHUTDOWN_META_KEY = "logurich_shutdown_registered" F = TypeVar("F", bound=Callable[..., Any]) @@ -134,6 +135,10 @@ def click_logger_init( level_by_module=lbm, rich_handler=logger_rich, ) + click_ctx = click.get_current_context(silent=True) + if click_ctx is not None and not click_ctx.meta.get(_CLICK_SHUTDOWN_META_KEY): + click_ctx.call_on_close(shutdown_logger) + click_ctx.meta[_CLICK_SHUTDOWN_META_KEY] = True logger.debug("Log level: %s", logger_level) logger.debug("Log verbose: %s", logger_verbose) logger.debug("Log filename: %s", logger_filename) diff --git a/src/logurich/struct.py b/src/logurich/struct.py index 140f55a..cf641ea 100644 --- a/src/logurich/struct.py +++ b/src/logurich/struct.py @@ -11,4 +11,5 @@ "final_handlers": (), "env_extra": {}, "atexit_registered": False, + "threading_atexit_registered": False, } diff --git a/tests/test_core.py b/tests/test_core.py index c0cbda6..0e2cce9 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -1,5 +1,6 @@ import json import logging +import threading from types import MappingProxyType import pytest @@ -34,11 +35,20 @@ def test_level_info(logger, buffer): assert "Debug, world!" not in output -def test_init_logger_registers_atexit_shutdown_once(monkeypatch): +def test_init_logger_registers_shutdown_hooks_once(monkeypatch): registered: list[object] = [] + thread_registered: list[object] = [] + register_threading_atexit = getattr(threading, "_register_atexit", None) monkeypatch.setitem(logger_state, "atexit_registered", False) + monkeypatch.setitem(logger_state, "threading_atexit_registered", False) monkeypatch.setattr("logurich.core.atexit.register", registered.append) + if register_threading_atexit is None: + monkeypatch.delattr("logurich.core.threading._register_atexit", raising=False) + else: + monkeypatch.setattr( + "logurich.core.threading._register_atexit", thread_registered.append + ) init_logger("INFO", enqueue=False) shutdown_logger() @@ -46,6 +56,10 @@ def test_init_logger_registers_atexit_shutdown_once(monkeypatch): shutdown_logger() assert registered == [shutdown_logger] + if register_threading_atexit is None: + assert thread_registered == [] + else: + assert thread_registered == [shutdown_logger] @pytest.mark.parametrize( diff --git a/tests/test_mp.py b/tests/test_mp.py index fbce7de..12f83ba 100644 --- a/tests/test_mp.py +++ b/tests/test_mp.py @@ -91,6 +91,49 @@ def test_rich_logging_in_child_process(buffer): assert "Rich Test" in output +def test_interpreter_exit_stops_queue_listener_without_thread_error(tmp_path): + repo_root = Path(__file__).resolve().parents[1] + script_path = tmp_path / "queue_listener_exit.py" + script_path.write_text( + textwrap.dedent( + """ + import logging + + from logurich import init_logger + + + def main(): + init_logger("INFO", enqueue=True) + logging.getLogger("exit").info("Interpreter exit test") + + + if __name__ == "__main__": + main() + """ + ) + ) + + env = os.environ.copy() + pythonpath = str(repo_root / "src") + if env.get("PYTHONPATH"): + pythonpath = f"{pythonpath}{os.pathsep}{env['PYTHONPATH']}" + env["PYTHONPATH"] = pythonpath + + result = subprocess.run( + [sys.executable, str(script_path)], + cwd=repo_root, + capture_output=True, + text=True, + env=env, + timeout=30, + ) + + assert result.returncode == 0, result.stderr or result.stdout + assert "Interpreter exit test" in result.stdout + assert "Exception in thread" not in result.stderr + assert "handle is closed" not in result.stderr + + def test_spawn_pool_initializer_can_configure_child_logging(tmp_path): repo_root = Path(__file__).resolve().parents[1] script_path = tmp_path / "spawn_pool_logging.py" diff --git a/tests/test_opt_click.py b/tests/test_opt_click.py index 81eb122..3baa2d4 100644 --- a/tests/test_opt_click.py +++ b/tests/test_opt_click.py @@ -1,9 +1,10 @@ import pytest click = pytest.importorskip("click") +from click.testing import CliRunner # noqa: E402 from logurich import shutdown_logger # noqa: E402 -from logurich.opt_click import click_logger_init # noqa: E402 +from logurich.opt_click import click_logger_init, click_logger_params # noqa: E402 from logurich.struct import logger_state # noqa: E402 @@ -16,4 +17,23 @@ def test_click_logger_init_registers_atexit_shutdown(monkeypatch): click_logger_init("INFO", 0, None, (), False) shutdown_logger() - assert registered == [shutdown_logger] + assert registered.count(shutdown_logger) == 1 + + +def test_click_logger_params_registers_context_shutdown(monkeypatch): + shutdown_calls: list[str] = [] + + monkeypatch.setattr("logurich.opt_click.init_logger", lambda *args, **kwargs: None) + monkeypatch.setattr( + "logurich.opt_click.shutdown_logger", lambda: shutdown_calls.append("called") + ) + + @click.command() + @click_logger_params + def cli() -> None: + return None + + result = CliRunner().invoke(cli, []) + + assert result.exit_code == 0 + assert shutdown_calls == ["called"]