Passing the method implementation to the formatter class

Author: leandrodamascena

aws_lambda_powertools/logging/formatter.py

@@ -7,10 +7,11 @@
 import time
 import traceback
 from abc import ABCMeta, abstractmethod
+from contextlib import contextmanager
 from contextvars import ContextVar
 from datetime import datetime, timezone
 from functools import partial
-from typing import TYPE_CHECKING, Any, Callable, Iterable
+from typing import TYPE_CHECKING, Any, Callable, Generator, Iterable
 
 from aws_lambda_powertools.shared import constants
 from aws_lambda_powertools.shared.functions import powertools_dev_is_set
@@ -62,6 +63,10 @@ def clear_state(self) -> None:
         """Removes any previously added logging keys"""
         raise NotImplementedError()
 
+    @contextmanager
+    def append_context_keys(self, **additional_keys: Any) -> Generator[None, None, None]:
+        yield
+
     # These specific thread-safe methods are necessary to manage shared context in concurrent environments.
     # They prevent race conditions and ensure data consistency across multiple threads.
     def thread_safe_append_keys(self, **additional_keys) -> None:
@@ -263,6 +268,31 @@ def clear_state(self) -> None:
         self.log_format = dict.fromkeys(self.log_record_order)
         self.log_format.update(**self.keys_combined)
 
+    @contextmanager
+    def append_context_keys(self, **additional_keys: Any) -> Generator[None, None, None]:
+        """
+        Context manager to temporarily add logging keys.
+
+        Parameters:
+        -----------
+        **keys: Any
+            Key-value pairs to include in the log context during the lifespan of the context manager.
+
+        Example:
+        --------
+        >>> logger = Logger(service="example_service")
+        >>> with logger.append_context_keys(user_id="123", operation="process"):
+        >>>     logger.info("Log with context")
+        >>> logger.info("Log without context")
+        """
+        # Add keys to the context
+        self.append_keys(**additional_keys)
+        try:
+            yield
+        finally:
+            # Remove the keys after exiting the context
+            self.remove_keys(additional_keys.keys())
+
     # These specific thread-safe methods are necessary to manage shared context in concurrent environments.
     # They prevent race conditions and ensure data consistency across multiple threads.
     def thread_safe_append_keys(self, **additional_keys) -> None:

aws_lambda_powertools/logging/logger.py

@@ -330,31 +330,6 @@ def _configure_sampling(self) -> None:
                 ),
             )
 
-    @contextmanager
-    def append_context_keys(self, **keys: Any) -> Generator[Any, Any, Any]:
-        """
-        Context manager to temporarily add logging keys.
-
-        Parameters:
-        -----------
-        **keys: Any
-            Key-value pairs to include in the log context during the lifespan of the context manager.
-
-        Example:
-        --------
-        >>> logger = Logger(service="example_service")
-        >>> with logger.append_context_keys(user_id="123", operation="process"):
-        >>>     logger.info("Log with context")
-        >>> logger.info("Log without context")
-        """
-        # Add keys to the context
-        self.append_keys(**keys)
-        try:
-            yield
-        finally:
-            # Remove the keys after exiting the context
-            self.remove_keys(keys.keys())
-
     @overload
     def inject_lambda_context(
         self,
@@ -606,6 +581,26 @@ def get_current_keys(self) -> dict[str, Any]:
     def remove_keys(self, keys: Iterable[str]) -> None:
         self.registered_formatter.remove_keys(keys)
 
+    @contextmanager
+    def append_context_keys(self, **additional_keys: Any) -> Generator[None, None, None]:
+        """
+        Context manager to temporarily add logging keys.
+
+        Parameters:
+        -----------
+        **keys: Any
+            Key-value pairs to include in the log context during the lifespan of the context manager.
+
+        Example:
+        --------
+        >>> logger = Logger(service="example_service")
+        >>> with logger.append_context_keys(user_id="123", operation="process"):
+        >>>     logger.info("Log with context")
+        >>> logger.info("Log without context")
+        """
+        with self.registered_formatter.append_context_keys(**additional_keys):
+            yield
+
     # These specific thread-safe methods are necessary to manage shared context in concurrent environments.
     # They prevent race conditions and ensure data consistency across multiple threads.
     def thread_safe_append_keys(self, **additional_keys: object) -> None: