feat(logger): type log record in LambdaPowertoolsFormatter with TypedDict (aws-powertools#2419)

Co-authored-by: erikayao93 <[email protected]>
Co-authored-by: Heitor Lessa <[email protected]>
Co-authored-by: Leandro Damascena <[email protected]>

Author: 3 people

.pre-commit-config.yaml

@@ -5,7 +5,7 @@
 
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v2.4.0
+    rev: "f71fa2c1f9cf5cb705f73dffe4b21f7c61470ba9" # v4.4.0
     hooks:
       - id: check-merge-conflict
       - id: trailing-whitespace
@@ -30,9 +30,9 @@ repos:
         language: system
         types: [python]
   - repo: https://github.com/igorshubovych/markdownlint-cli
-    rev: "11c08644ce6df850480d98f628596446a526cbc6" # frozen: v0.31.1
+    rev: "ce0d77ac47dc921b62429804fe763d4d35f66a76" # v0.34.0
     hooks:
-      - id: markdownlint
+      - id: markdownlint-docker
         args: ["--fix"]
   - repo: local
     hooks:
@@ -43,7 +43,7 @@ repos:
         types: [yaml]
         files: examples/.*
   - repo: https://github.com/rhysd/actionlint
-    rev: v1.6.23
+    rev: "fd7ba3c382e13dcc0248e425b4cbc3f1185fa3ee" # v1.6.24
     hooks:
       - id: actionlint-docker
         args: [-pyflakes=]

aws_lambda_powertools/logging/formatter.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import inspect
 import json
 import logging
@@ -8,8 +10,9 @@
 from functools import partial
 from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
 
-from ..shared import constants
-from ..shared.functions import powertools_dev_is_set
+from aws_lambda_powertools.logging.types import LogRecord
+from aws_lambda_powertools.shared import constants
+from aws_lambda_powertools.shared.functions import powertools_dev_is_set
 
 RESERVED_LOG_ATTRS = (
     "name",
@@ -66,12 +69,12 @@ class LambdaPowertoolsFormatter(BasePowertoolsFormatter):
 
     def __init__(
         self,
-        json_serializer: Optional[Callable[[Dict], str]] = None,
-        json_deserializer: Optional[Callable[[Union[Dict, str, bool, int, float]], str]] = None,
-        json_default: Optional[Callable[[Any], Any]] = None,
-        datefmt: Optional[str] = None,
+        json_serializer: Callable[[LogRecord], str] | None = None,
+        json_deserializer: Callable[[Dict | str | bool | int | float], str] | None = None,
+        json_default: Callable[[Any], Any] | None = None,
+        datefmt: str | None = None,
         use_datetime_directive: bool = False,
-        log_record_order: Optional[List[str]] = None,
+        log_record_order: List[str] | None = None,
         utc: bool = False,
         use_rfc3339: bool = False,
         **kwargs,
@@ -144,7 +147,7 @@ def __init__(
 
         super().__init__(datefmt=self.datefmt)
 
-    def serialize(self, log: Dict) -> str:
+    def serialize(self, log: LogRecord) -> str:
         """Serialize structured log dict to JSON str"""
         return self.json_serializer(log)
 

aws_lambda_powertools/logging/formatters/datadog.py

@@ -1,15 +1,16 @@
 from __future__ import annotations
 
-from typing import Any, Callable
+from typing import Any, Callable, Dict
 
 from aws_lambda_powertools.logging.formatter import LambdaPowertoolsFormatter
+from aws_lambda_powertools.logging.types import LogRecord
 
 
 class DatadogLogFormatter(LambdaPowertoolsFormatter):
     def __init__(
         self,
-        json_serializer: Callable[[dict], str] | None = None,
-        json_deserializer: Callable[[dict | str | bool | int | float], str] | None = None,
+        json_serializer: Callable[[LogRecord], str] | None = None,
+        json_deserializer: Callable[[Dict | str | bool | int | float], str] | None = None,
         json_default: Callable[[Any], Any] | None = None,
         datefmt: str | None = None,
         use_datetime_directive: bool = False,

aws_lambda_powertools/logging/types.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+import sys
+
+if sys.version_info >= (3, 11):
+    from typing import NotRequired
+else:
+    from typing_extensions import NotRequired
+
+if sys.version_info >= (3, 8):
+    from typing import TypedDict
+else:
+    from typing_extensions import TypedDict
+
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias
+
+from typing import Any, Dict, List, Union
+
+LogRecord: TypeAlias = Union[Dict[str, Any], "PowertoolsLogRecord"]
+
+
+class PowertoolsLogRecord(TypedDict):
+    # Base fields (required)
+    level: str
+    location: str
+    message: Dict[str, Any] | str | bool | List[Any]
+    timestamp: str | int
+    service: str
+
+    # Fields from logger.inject_lambda_context
+    cold_start: NotRequired[bool]
+    function_name: NotRequired[str]
+    function_memory_size: NotRequired[int]
+    function_arn: NotRequired[str]
+    function_request_id: NotRequired[str]
+    # From logger.inject_lambda_context if AWS X-Ray is enabled
+    xray_trace_id: NotRequired[str]
+
+    # If sample_rate is defined
+    sampling_rate: NotRequired[float]
+
+    # From logger.set_correlation_id
+    correlation_id: NotRequired[str]
+
+    # Fields from logger.exception
+    exception_name: NotRequired[str]
+    exception: NotRequired[str]

docs/core/logger.md

@@ -604,7 +604,7 @@ For these, you can override the `serialize` method from [LambdaPowertoolsFormatt
 
 === "bring_your_own_formatter.py"
 
-    ```python hl_lines="2 5-6 12"
+    ```python hl_lines="2-3 6 11-12 15"
     --8<-- "examples/logger/src/bring_your_own_formatter.py"
     ```
 

examples/logger/src/bring_your_own_formatter.py

@@ -1,12 +1,15 @@
 from aws_lambda_powertools import Logger
 from aws_lambda_powertools.logging.formatter import LambdaPowertoolsFormatter
+from aws_lambda_powertools.logging.types import LogRecord
 
 
 class CustomFormatter(LambdaPowertoolsFormatter):
-    def serialize(self, log: dict) -> str:
+    def serialize(self, log: LogRecord) -> str:
         """Serialize final structured log dict to JSON str"""
-        log["event"] = log.pop("message")  # rename message key to event
-        return self.json_serializer(log)  # use configured json serializer
+        # in this example, log["message"] is a required field
+        # but we want to remap to "event" and delete "message", hence mypy ignore checks
+        log["event"] = log.pop("message")  # type: ignore[typeddict-unknown-key,misc]
+        return self.json_serializer(log)
 
 
 logger = Logger(service="payment", logger_formatter=CustomFormatter())