From bebd89f767c39747d3731f8f2bb307a316a9d1fe Mon Sep 17 00:00:00 2001
From: James Mishra <james.mishra@neocrym.com>
Date: Mon, 23 Oct 2023 16:27:10 -0700
Subject: [PATCH] Give the ability to set `extra` when initializing a logger

---
 log_with_context/__init__.py | 22 +++++++++++++++++++---
 test_log_with_context.py     | 28 ++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/log_with_context/__init__.py b/log_with_context/__init__.py
index 6650f98..41d9754 100644
--- a/log_with_context/__init__.py
+++ b/log_with_context/__init__.py
@@ -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
@@ -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."""
diff --git a/test_log_with_context.py b/test_log_with_context.py
index 59c89ac..68383ce 100644
--- a/test_log_with_context.py
+++ b/test_log_with_context.py
@@ -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
+                )