chore(common): reusable function to extract event from models

Author: heitorlessa

aws_lambda_powertools/shared/functions.py

@@ -1,10 +1,13 @@
+from __future__ import annotations
+
 import base64
+import dataclasses
 import itertools
 import logging
 import os
 import warnings
 from binascii import Error as BinAsciiError
-from typing import Dict, Generator, Optional, Union, overload
+from typing import Any, Dict, Generator, Optional, Union, overload
 
 from aws_lambda_powertools.shared import constants
 
@@ -121,3 +124,43 @@ def powertools_debug_is_set() -> bool:
 def slice_dictionary(data: Dict, chunk_size: int) -> Generator[Dict, None, None]:
     for _ in range(0, len(data), chunk_size):
         yield {dict_key: data[dict_key] for dict_key in itertools.islice(data, chunk_size)}
+
+
+def extract_event_from_common_models(data: Any) -> Dict | Any:
+    """Extract raw event from common types used in Powertools
+
+    If event cannot be extracted, return received data as is.
+
+    Common models:
+
+        - Event Source Data Classes (DictWrapper)
+        - Python Dataclasses
+        - Pydantic Models (BaseModel)
+
+    Parameters
+    ----------
+    data : Any
+        Original event, a potential instance of DictWrapper/BaseModel/Dataclass
+
+    Notes
+    -----
+
+    Why not using static type for function argument?
+
+    DictWrapper would cause a circular import. Pydantic BaseModel could
+    cause a ModuleNotFound or trigger init reflection worsening cold start.
+    """
+    # Short-circuit most common type first for perf
+    if isinstance(data, dict):
+        return data
+
+    # Is it an Event Source Data Class?
+    if getattr(data, "raw_event", None):
+        return data.raw_event
+
+    # Is it a Pydantic Model?
+    if callable(getattr(data, "dict", None)):
+        return data.dict()
+
+    # Is it a Dataclass? If not return as is
+    return dataclasses.asdict(data) if dataclasses.is_dataclass(data) else data

tests/functional/test_shared_functions.py

@@ -1,15 +1,19 @@
 import warnings
+from dataclasses import dataclass
 
 import pytest
+from pydantic import BaseModel
 
 from aws_lambda_powertools.shared import constants
 from aws_lambda_powertools.shared.functions import (
+    extract_event_from_common_models,
     powertools_debug_is_set,
     powertools_dev_is_set,
     resolve_env_var_choice,
     resolve_truthy_env_var_choice,
     strtobool,
 )
+from aws_lambda_powertools.utilities.data_classes.common import DictWrapper
 
 
 def test_resolve_env_var_choice_explicit_wins_over_env_var():
@@ -64,3 +68,38 @@ def test_powertools_debug_warning(monkeypatch: pytest.MonkeyPatch):
         powertools_debug_is_set()
         assert len(w) == 1
         assert str(w[0].message) == warning_message
+
+
+def test_extract_event_dict():
+    data = {"hello": "world"}
+    assert extract_event_from_common_models(data) == data
+
+
+def test_extract_event_pydantic():
+    class DummyModel(BaseModel):
+        hello: str
+
+    data = {"hello": "world"}
+    assert extract_event_from_common_models(DummyModel(**data)) == data
+
+
+def test_extract_event_dict_wrapper():
+    class DummyClassSample(DictWrapper):
+        pass
+
+    data = {"hello": "world"}
+    assert extract_event_from_common_models(DummyClassSample(data)) == data
+
+
+def test_extract_event_dataclass():
+    @dataclass
+    class DummyDataclass:
+        hello: str
+
+    data = {"hello": "world"}
+    assert extract_event_from_common_models(DummyDataclass(**data)) == data
+
+
+@pytest.mark.parametrize("data", [False, True, "", 10, [], {}, object])
+def test_extract_event_any(data):
+    assert extract_event_from_common_models(data) == data