From a53e86dd41eb748613a3d74e2c024a2592a1dbb7 Mon Sep 17 00:00:00 2001 From: Terseus Date: Thu, 6 Feb 2025 19:50:04 +0100 Subject: [PATCH 1/3] test: add multithreading test --- src/logging_with_context/global_context.py | 10 +++--- .../test_global_context.py | 34 +++++++++++++++++++ 2 files changed, 40 insertions(+), 4 deletions(-) diff --git a/src/logging_with_context/global_context.py b/src/logging_with_context/global_context.py index a28b931..735acb8 100644 --- a/src/logging_with_context/global_context.py +++ b/src/logging_with_context/global_context.py @@ -18,6 +18,10 @@ ) +def _get_loggers_to_process(loggers: Optional[Sequence[Logger]] = None) -> list[Logger]: + return [getLogger()] if loggers is None else list(loggers) + + def init_global_context(loggers: Optional[Sequence[Logger]] = None) -> None: """ Initialize the application global context in the given loggers. @@ -26,9 +30,8 @@ def init_global_context(loggers: Optional[Sequence[Logger]] = None) -> None: loggers: The loggers to attach the global context; if not loggers are specified it will use the root logger. """ - loggers_to_process = [getLogger()] if loggers is None else list(loggers) filter_with_context = FilterWithContextVar(__global_context_var) - for logger in loggers_to_process: + for logger in _get_loggers_to_process(loggers): for handler in logger.handlers: handler.addFilter(filter_with_context) @@ -41,8 +44,7 @@ def shutdown_global_context(loggers: Optional[Sequence[Logger]] = None) -> None: loggers: The loggers that were used when calling `init_global_context`; by default the root logger. """ - loggers_to_process = [getLogger()] if loggers is None else list(loggers) - for logger in loggers_to_process: + for logger in _get_loggers_to_process(loggers): for handler in logger.handlers: for filter_ in handler.filters: if not isinstance(filter_, FilterWithContextVar): diff --git a/tests/logging_with_context/test_global_context.py b/tests/logging_with_context/test_global_context.py index 9e4c9b7..fd78306 100644 --- a/tests/logging_with_context/test_global_context.py +++ b/tests/logging_with_context/test_global_context.py @@ -1,4 +1,5 @@ import logging +import threading import pytest @@ -40,3 +41,36 @@ def test_add_global_context_after_shutdown_ignored_ok(caplog: pytest.LogCaptureF assert len(caplog.records) == 1 result = caplog.records[0] assert not hasattr(result, "key") + + +def test_add_global_context_multithread(caplog: pytest.LogCaptureFixture): + def worker(value: int) -> None: + with add_global_context({"value": value}): + logger.info("Message 1 from thread %s", value) + with add_global_context({"value": value * 10}): + logger.info("Message 2 from thread %s", value) + with add_global_context({"value": value * 100}): + logger.info("Message 3 from thread %s", value) + + logger = logging.getLogger(__name__) + with global_context_initialized(), caplog.at_level(logging.INFO): + worker_1 = threading.Thread(target=worker, args=(1,)) + worker_2 = threading.Thread(target=worker, args=(2,)) + worker_1.start() + worker_2.start() + worker_1.join() + worker_2.join() + assert len(caplog.records) == 6 + result = [ + {"message": record.message, "value": record.value} # type: ignore + for record in sorted(caplog.records, key=lambda r: r.value) # type: ignore + ] + expected = [ + {"message": "Message 1 from thread 1", "value": 1}, + {"message": "Message 1 from thread 2", "value": 2}, + {"message": "Message 2 from thread 1", "value": 10}, + {"message": "Message 2 from thread 2", "value": 20}, + {"message": "Message 3 from thread 1", "value": 100}, + {"message": "Message 3 from thread 2", "value": 200}, + ] + assert result == expected From 434372bd2851ab8bf322c585a5103ac1ff86c4f6 Mon Sep 17 00:00:00 2001 From: Terseus Date: Thu, 6 Feb 2025 19:50:44 +0100 Subject: [PATCH 2/3] ci: replace --locked with --frozen Otherwise sometimes `uv` fails because changes due to not being in a version tag. --- run_tests.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/run_tests.sh b/run_tests.sh index 85759f1..3c61a73 100755 --- a/run_tests.sh +++ b/run_tests.sh @@ -1,5 +1,5 @@ #!/usr/bin/env bash for python in 3.9 3.10 3.11 3.12 3.13; do - uv run --locked --isolated --python=$python pytest + uv run --frozen --isolated --python=$python pytest done From d4c18287b4c49803186e8ec5c10252c13970c4b4 Mon Sep 17 00:00:00 2001 From: Terseus Date: Thu, 6 Feb 2025 20:29:08 +0100 Subject: [PATCH 3/3] feat: add auto_init to add_global_context --- docs/howto.md | 20 +++++++++++++++- src/logging_with_context/global_context.py | 23 ++++++++++++++++++- .../test_global_context.py | 16 +++++++++++-- 3 files changed, 55 insertions(+), 4 deletions(-) diff --git a/docs/howto.md b/docs/howto.md index 492feef..bf15a1e 100644 --- a/docs/howto.md +++ b/docs/howto.md @@ -45,9 +45,27 @@ with add_global_context({"user_id": user_id, "request_id": request_id}): ``` +## Using the automatic init + +If you don't want to customise the initialization you can let `add_global_context` automatically handle the init and shutdown for you: + +```python +import logging + +from logging_with_context.global_context import add_global_context + + +def main(): + logging.basicConfig(level=logging.INFO) # Or any other way to setup logging. + with add_global_context({"user_id": 10}): + # Here the context is automatically initialized. + # It'll also be automatically shutdown once this context manager finishes. +``` + + ## Using the init/shutdown API -In case you can't use the context manager, you can use the manual initialization and shutdown API: +In case you want to customise the initialization but can't use the context manager, you can use the manual initialization and shutdown API: ```python import logging diff --git a/src/logging_with_context/global_context.py b/src/logging_with_context/global_context.py index 735acb8..f03c29c 100644 --- a/src/logging_with_context/global_context.py +++ b/src/logging_with_context/global_context.py @@ -16,6 +16,9 @@ __global_context_var: ContextVar[dict[str, Any]] = ContextVar( "global_context", default={} ) +__global_context_initialized: ContextVar[bool] = ContextVar( + "global_context_initialized", default=False +) def _get_loggers_to_process(loggers: Optional[Sequence[Logger]] = None) -> list[Logger]: @@ -30,6 +33,7 @@ def init_global_context(loggers: Optional[Sequence[Logger]] = None) -> None: loggers: The loggers to attach the global context; if not loggers are specified it will use the root logger. """ + __global_context_initialized.set(True) filter_with_context = FilterWithContextVar(__global_context_var) for logger in _get_loggers_to_process(loggers): for handler in logger.handlers: @@ -44,6 +48,7 @@ def shutdown_global_context(loggers: Optional[Sequence[Logger]] = None) -> None: loggers: The loggers that were used when calling `init_global_context`; by default the root logger. """ + __global_context_initialized.set(False) for logger in _get_loggers_to_process(loggers): for handler in logger.handlers: for filter_ in handler.filters: @@ -74,7 +79,9 @@ def global_context_initialized( @contextmanager -def add_global_context(context: dict[str, Any]) -> Generator[None, None, None]: +def add_global_context( + context: dict[str, Any], *, auto_init: bool = True +) -> Generator[None, None, None]: """ Add values to the global context to be attached to all the log messages. @@ -82,12 +89,26 @@ def add_global_context(context: dict[str, Any]) -> Generator[None, None, None]: Parameters: context: A key/value mapping with the values to add to the global context. + auto_init: Indicate if the global context should be automatically initialized + if it isn't. + + If `True`, the context will be also automatically shutdown before exiting. + + If the global context is already initialized it'll do nothing. + + Keyword-only argument. Returns: A context manager that manages the life of the values. """ + auto_initialized = False + if not __global_context_initialized.get() and auto_init: + init_global_context() + auto_initialized = True token = __global_context_var.set(__global_context_var.get() | context) try: yield finally: __global_context_var.reset(token) + if auto_initialized: + shutdown_global_context() diff --git a/tests/logging_with_context/test_global_context.py b/tests/logging_with_context/test_global_context.py index fd78306..3370695 100644 --- a/tests/logging_with_context/test_global_context.py +++ b/tests/logging_with_context/test_global_context.py @@ -24,7 +24,7 @@ def test_add_global_context_ok(caplog: pytest.LogCaptureFixture): def test_add_global_context_without_init_ignored_ok(caplog: pytest.LogCaptureFixture): logger = logging.getLogger(__name__) - with add_global_context({"key": "value"}): + with add_global_context({"key": "value"}, auto_init=False): with caplog.at_level(logging.INFO): logger.info("Test message") assert len(caplog.records) == 1 @@ -36,13 +36,25 @@ def test_add_global_context_after_shutdown_ignored_ok(caplog: pytest.LogCaptureF logger = logging.getLogger(__name__) with global_context_initialized(): pass - with add_global_context({"key": "value"}), caplog.at_level(logging.INFO): + with ( + add_global_context({"key": "value"}, auto_init=False), + caplog.at_level(logging.INFO), + ): logger.info("Test message") assert len(caplog.records) == 1 result = caplog.records[0] assert not hasattr(result, "key") +def test_add_global_context_auto_init_ok(caplog: pytest.LogCaptureFixture): + logger = logging.getLogger(__name__) + with add_global_context({"key": "value"}), caplog.at_level(logging.INFO): + logger.info("Test message") + assert len(caplog.records) == 1 + result = caplog.records[0] + assert result.key == "value" # type: ignore + + def test_add_global_context_multithread(caplog: pytest.LogCaptureFixture): def worker(value: int) -> None: with add_global_context({"value": value}):