Skip to content

Commit

Permalink
Give the ability to set extra when initializing a logger
Browse files Browse the repository at this point in the history
  • Loading branch information
jamesmishra committed Oct 23, 2023
1 parent e9b9aae commit bebd89f
Show file tree
Hide file tree
Showing 2 changed files with 47 additions and 3 deletions.
22 changes: 19 additions & 3 deletions log_with_context/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -61,18 +61,31 @@ class Logger:
def __init__(
self,
name: Optional[str] = None,
logger: Optional[Union[logging.Logger, "Logger"]] = None,
):
logger: Union[logging.Logger, "Logger", None] = None,
*,
initial_extra: Optional[Mapping[str, Any]] = None,
) -> None:
"""
Initialize.
Args:
name: The name of the logger.
logger: A logging.Logger instance to wrap.
initial_extra: An optional mapping of key/value pairs to
initialize the logger with. The values in `initial_extra`
are specific to this logger, and not to other loggers
in the same process.
"""
if isinstance(logger, Logger):
self.base_logger: Any = logger.base_logger
else:
self.base_logger = logger or logging.getLogger(name=name)
if initial_extra:
self._initial_extra = initial_extra
else:
self._initial_extra = {}

def _msg(
self, func: Callable[..., None], msg: Any, *args: Any, **kwargs: Any
Expand All @@ -92,7 +105,10 @@ def _msg(
@property
def extra(self) -> _EXTRA_TYPE:
"""Return the extra metadata that this logger sends with every message."""
return get_extra()
new_extra = get_extra()
if self._initial_extra:
return {**self._initial_extra, **new_extra}
return new_extra

def debug(self, msg: Any, *args: Any, **kwargs: Any) -> None:
"""Debug."""
Expand Down
28 changes: 28 additions & 0 deletions test_log_with_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -119,3 +119,31 @@ def test_process_local(self) -> None:
with concurrent.futures.ProcessPoolExecutor(max_workers=2) as exc:
with add_logging_context(a=1):
list(exc.map(process_worker, [1, 2]))

def test_extra_without_add_logging_context(self) -> None:
"""
Test that we can assign an initial extra value
when we initialize the logger.
"""
base_logger = unittest.mock.Mock(spec=LOGGER)
logger = Logger(logger=base_logger, initial_extra=dict(a=1))
self.assertEqual(logger.extra, dict(a=1))
self.assertEqual(self.logger.extra, dict())
logger.debug("1")
base_logger.debug.assert_called_with("1", extra=dict(a=1), stacklevel=3)
with add_logging_context(b=2):
self.assertEqual(logger.extra, dict(a=1, b=2))
self.assertEqual(self.logger.extra, dict(b=2))
logger.info("2")
base_logger.info.assert_called_with("2", extra=dict(a=1, b=2), stacklevel=3)
with add_logging_context(c=3):
self.assertEqual(logger.extra, dict(a=1, b=2, c=3))
self.assertEqual(self.logger.extra, dict(b=2, c=3))
logger.info("3")
base_logger.info.assert_called_with(
"3", extra=dict(a=1, b=2, c=3), stacklevel=3
)
logger.info("4", extra=dict(d=4))
base_logger.info.assert_called_with(
"4", extra=dict(a=1, b=2, c=3, d=4), stacklevel=3
)

0 comments on commit bebd89f

Please sign in to comment.