Apply PR suggestions

Author: Michal Ploski

aws_lambda_powertools/utilities/batch/sqs.py

@@ -90,10 +90,13 @@ def _get_entries_to_clean(self) -> List[Dict[str, str]]:
         """
         Format messages to use in batch deletion
         """
-        return [
-            {"Id": msg["messageId"], "ReceiptHandle": msg["receiptHandle"]}
-            for msg in cast(List[SQSRecord], self.success_messages)
-        ]
+        entries = []
+        # success_messages has generic type of union of SQS, Dynamodb and Kinesis Streams records or Pydantic models.
+        # Here we get SQS Record only
+        messages = cast(List[SQSRecord], self.success_messages)
+        for msg in messages:
+            entries.append({"Id": msg["messageId"], "ReceiptHandle": msg["receiptHandle"]})
+        return entries
 
     def _process_record(self, record) -> Tuple:
         """

aws_lambda_powertools/utilities/idempotency/persistence/base.py

@@ -40,7 +40,7 @@ def __init__(
         idempotency_key,
         status: str = "",
         expiry_timestamp: Optional[int] = None,
-        response_data: str = "",
+        response_data: Optional[str] = "",
         payload_hash: Optional[str] = None,
     ) -> None:
         """
@@ -92,16 +92,16 @@ def status(self) -> str:
         else:
             raise IdempotencyInvalidStatusError(self._status)
 
-    def response_json_as_dict(self) -> dict:
+    def response_json_as_dict(self) -> Optional[dict]:
         """
         Get response data deserialized to python dict
 
         Returns
         -------
-        dict
+        Optional[dict]
             previous response data deserialized
         """
-        return json.loads(self.response_data)
+        return json.loads(self.response_data) if self.response_data else None
 
 
 class BasePersistenceLayer(ABC):
@@ -121,7 +121,6 @@ def __init__(self):
         self.raise_on_no_idempotency_key = False
         self.expires_after_seconds: int = 60 * 60  # 1 hour default
         self.use_local_cache = False
-        self._cache: Optional[LRUDict] = None
         self.hash_function = None
 
     def configure(self, config: IdempotencyConfig, function_name: Optional[str] = None) -> None:
@@ -279,14 +278,14 @@ def _save_to_cache(self, data_record: DataRecord):
         -------
 
         """
-        if not self.use_local_cache or self._cache is None:
+        if not self.use_local_cache:
             return
         if data_record.status == STATUS_CONSTANTS["INPROGRESS"]:
             return
         self._cache[data_record.idempotency_key] = data_record
 
     def _retrieve_from_cache(self, idempotency_key: str):
-        if not self.use_local_cache or self._cache is None:
+        if not self.use_local_cache:
             return
         cached_record = self._cache.get(key=idempotency_key)
         if cached_record:
@@ -296,7 +295,7 @@ def _retrieve_from_cache(self, idempotency_key: str):
             self._delete_from_cache(idempotency_key=idempotency_key)
 
     def _delete_from_cache(self, idempotency_key: str):
-        if not self.use_local_cache or self._cache is None:
+        if not self.use_local_cache:
             return
         if idempotency_key in self._cache:
             del self._cache[idempotency_key]

aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py

@@ -130,7 +130,7 @@ def _item_to_data_record(self, item: Dict[str, Any]) -> DataRecord:
             idempotency_key=item[self.key_attr],
             status=item[self.status_attr],
             expiry_timestamp=item[self.expiry_attr],
-            response_data=item.get(self.data_attr, ""),
+            response_data=item.get(self.data_attr),
             payload_hash=item.get(self.validation_key_attr),
         )
 

aws_lambda_powertools/utilities/parameters/ssm.py

@@ -87,6 +87,8 @@ def __init__(self, config: Optional[Config] = None, boto3_session: Optional[boto
 
         super().__init__()
 
+    # We break Liskov substitution principle due to differences in signatures of this method and superclass get method
+    # We ignore mypy error, as changes to the signature here or in a superclass is a breaking change to users
     def get(  # type: ignore[override]
         self,
         name: str,

aws_lambda_powertools/utilities/parser/envelopes/kinesis.py

@@ -37,7 +37,9 @@ def parse(self, data: Optional[Union[Dict[str, Any], Any]], model: Type[Model])
         logger.debug(f"Parsing incoming data with Kinesis model {KinesisDataStreamModel}")
         parsed_envelope: KinesisDataStreamModel = KinesisDataStreamModel.parse_obj(data)
         logger.debug(f"Parsing Kinesis records in `body` with {model}")
-        return [
-            self._parse(data=cast(bytes, record.kinesis.data).decode("utf-8"), model=model)
-            for record in parsed_envelope.Records
-        ]
+        models = []
+        for record in parsed_envelope.Records:
+            # We allow either AWS expected contract (bytes) or a custom Model, see #943
+            data = cast(bytes, record.kinesis.data)
+            models.append(self._parse(data=cast(bytes, record.kinesis.data).decode("utf-8"), model=model))
+        return models

aws_lambda_powertools/utilities/parser/envelopes/sns.py

@@ -69,6 +69,8 @@ def parse(self, data: Optional[Union[Dict[str, Any], Any]], model: Type[Model])
         parsed_envelope = SqsModel.parse_obj(data)
         output = []
         for record in parsed_envelope.Records:
-            sns_notification: SnsNotificationModel = SnsNotificationModel.parse_raw(cast(str, record.body))
+            # We allow either AWS expected contract (str) or a custom Model, see #943
+            body = cast(str, record.body)
+            sns_notification = SnsNotificationModel.parse_raw(body)
             output.append(self._parse(data=sns_notification.Message, model=model))
         return output

mypy.ini

@@ -31,3 +31,6 @@ ignore_missing_imports = True
 
 [mypy-aws_xray_sdk.ext.aiohttp.client]
 ignore_missing_imports = True
+
+[mypy-dataclasses]
+ignore_missing_imports = True

poetry.lock

@@ -54,7 +54,6 @@ mkdocs-material = "^7.3.6"
 mkdocs-git-revision-date-plugin = "^0.3.1"
 mike = "^0.6.0"
 mypy = "^0.931"
-types-dataclasses = "^0.6.4"
 
 
 [tool.poetry.extras]