fix(event_sources): implement Mapping protocol on DictWrapper for better interop with existing middlewares (#1516)

Co-authored-by: Heitor Lessa <[email protected]>
Co-authored-by: Ruben Fonseca <[email protected]>

Author: 3 people

aws_lambda_powertools/utilities/data_classes/common.py

@@ -1,9 +1,10 @@
 import base64
 import json
-from typing import Any, Dict, Optional
+from collections.abc import Mapping
+from typing import Any, Dict, Iterator, Optional
 
 
-class DictWrapper:
+class DictWrapper(Mapping):
     """Provides a single read only access to a wrapper dict"""
 
     def __init__(self, data: Dict[str, Any]):
@@ -19,6 +20,12 @@ def __eq__(self, other: Any) -> bool:
 
         return self._data == other._data
 
+    def __iter__(self) -> Iterator:
+        return iter(self._data)
+
+    def __len__(self) -> int:
+        return len(self._data)
+
     def get(self, key: str, default: Optional[Any] = None) -> Optional[Any]:
         return self._data.get(key, default)
 

aws_lambda_powertools/utilities/data_classes/dynamo_db_stream_event.py

@@ -182,8 +182,10 @@ def approximate_creation_date_time(self) -> Optional[int]:
         item = self.get("ApproximateCreationDateTime")
         return None if item is None else int(item)
 
+    # NOTE: This override breaks the Mapping protocol of DictWrapper, it's left here for backwards compatibility with
+    # a 'type: ignore' comment. See #1516 for discussion
     @property
-    def keys(self) -> Optional[Dict[str, AttributeValue]]:
+    def keys(self) -> Optional[Dict[str, AttributeValue]]:  # type: ignore[override]
         """The primary key attribute(s) for the DynamoDB item that was modified."""
         return _attribute_value_dict(self._data, "Keys")
 

tests/functional/test_data_classes.py

@@ -74,6 +74,7 @@
     AttributeValueType,
     DynamoDBRecordEventName,
     DynamoDBStreamEvent,
+    StreamRecord,
     StreamViewType,
 )
 from aws_lambda_powertools.utilities.data_classes.event_source import event_source
@@ -101,6 +102,19 @@ def message(self) -> str:
     assert DataClassSample(data1).raw_event is data1
 
 
+def test_dict_wrapper_implements_mapping():
+    class DataClassSample(DictWrapper):
+        pass
+
+    data = {"message": "foo1"}
+    event_source = DataClassSample(data)
+    assert len(event_source) == len(data)
+    assert list(event_source) == list(data)
+    assert event_source.keys() == data.keys()
+    assert list(event_source.values()) == list(data.values())
+    assert event_source.items() == data.items()
+
+
 def test_cloud_watch_dashboard_event():
     event = CloudWatchDashboardCustomWidgetEvent(load_event("cloudWatchDashboardEvent.json"))
     assert event.describe is False
@@ -617,6 +631,23 @@ def test_dynamo_attribute_value_type_error():
         print(attribute_value.get_type)
 
 
+def test_stream_record_keys_with_valid_keys():
+    attribute_value = {"Foo": "Bar"}
+    record = StreamRecord({"Keys": {"Key1": attribute_value}})
+    assert record.keys == {"Key1": AttributeValue(attribute_value)}
+
+
+def test_stream_record_keys_with_no_keys():
+    record = StreamRecord({})
+    assert record.keys is None
+
+
+def test_stream_record_keys_overrides_dict_wrapper_keys():
+    data = {"Keys": {"key1": {"attr1": "value1"}}}
+    record = StreamRecord(data)
+    assert record.keys != data.keys()
+
+
 def test_event_bridge_event():
     event = EventBridgeEvent(load_event("eventBridgeEvent.json"))