Fix mypy stuff

Author: leandrodamascena

aws_lambda_powertools/utilities/data_classes/kafka_event.py

@@ -10,7 +10,7 @@
     from collections.abc import Iterator
 
 
-class KafkaEventBase(DictWrapper):
+class KafkaEventRecordBase(DictWrapper):
     @property
     def topic(self) -> str:
         """The Kafka topic."""
@@ -37,7 +37,7 @@ def timestamp_type(self) -> str:
         return self["timestampType"]
 
 
-class KafkaEventRecord(KafkaEventBase):
+class KafkaEventRecord(KafkaEventRecordBase):
     @property
     def key(self) -> str | None:
         """
@@ -85,18 +85,7 @@ def decoded_headers(self) -> dict[str, bytes]:
         return CaseInsensitiveDict((k, bytes(v)) for chunk in self.headers for k, v in chunk.items())
 
 
-class KafkaEvent(DictWrapper):
-    """Self-managed or MSK Apache Kafka event trigger
-    Documentation:
-    --------------
-    - https://docs.aws.amazon.com/lambda/latest/dg/with-kafka.html
-    - https://docs.aws.amazon.com/lambda/latest/dg/with-msk.html
-    """
-
-    def __init__(self, data: dict[str, Any]):
-        super().__init__(data)
-        self._records: Iterator[KafkaEventRecord] | None = None
-
+class KafkaEventBase(DictWrapper):
     @property
     def event_source(self) -> str:
         """The AWS service from which the Kafka event record originated."""
@@ -117,6 +106,19 @@ def decoded_bootstrap_servers(self) -> list[str]:
         """The decoded Kafka bootstrap URL."""
         return self.bootstrap_servers.split(",")
 
+
+class KafkaEvent(KafkaEventBase):
+    """Self-managed or MSK Apache Kafka event trigger
+    Documentation:
+    --------------
+    - https://docs.aws.amazon.com/lambda/latest/dg/with-kafka.html
+    - https://docs.aws.amazon.com/lambda/latest/dg/with-msk.html
+    """
+
+    def __init__(self, data: dict[str, Any]):
+        super().__init__(data)
+        self._records: Iterator[KafkaEventRecord] | None = None
+
     @property
     def records(self) -> Iterator[KafkaEventRecord]:
         """The Kafka records."""

aws_lambda_powertools/utilities/kafka_consumer/consumer_records.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
+from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
 from aws_lambda_powertools.utilities.data_classes.common import CaseInsensitiveDict
-from aws_lambda_powertools.utilities.data_classes.kafka_event import KafkaEvent, KafkaEventBase
+from aws_lambda_powertools.utilities.data_classes.kafka_event import KafkaEventBase, KafkaEventRecordBase
 from aws_lambda_powertools.utilities.kafka_consumer.deserializer.deserializer import get_deserializer
 from aws_lambda_powertools.utilities.kafka_consumer.serialization.serialization import serialize_to_output_type
 
@@ -13,49 +14,49 @@
     from aws_lambda_powertools.utilities.kafka_consumer.schema_config import SchemaConfig
 
 
-class ConsumerRecordRecords(KafkaEventBase):
+class ConsumerRecordRecords(KafkaEventRecordBase):
     """
     A Kafka Consumer Record
     """
 
-    def __init__(self, data: dict[str, Any], deserialize: SchemaConfig | None = None):
+    def __init__(self, data: dict[str, Any], schema_config: SchemaConfig | None = None):
         super().__init__(data)
-        self.deserialize = deserialize
+        self.schema_config = schema_config
 
-    @property
+    @cached_property
     def key(self) -> Any:
         key = self.get("key")
-        if key and (self.deserialize and self.deserialize.key_schema_type):
+        if key and (self.schema_config and self.schema_config.key_schema_type):
             deserializer = get_deserializer(
-                self.deserialize.key_schema_type,
-                self.deserialize.key_schema_str,
+                self.schema_config.key_schema_type,
+                self.schema_config.key_schema_str,
             )
             deserialized_key = deserializer.deserialize(key)
 
-            if self.deserialize.key_output_serializer:
+            if self.schema_config.key_output_serializer:
                 return serialize_to_output_type(
                     deserialized_key,
-                    self.deserialize.key_output_serializer,
+                    self.schema_config.key_output_serializer,
                 )
 
             return deserialized_key
 
         return key
 
-    @property
+    @cached_property
     def value(self) -> Any:
         value = self["value"]
-        if value and (self.deserialize and self.deserialize.value_schema_type):
+        if value and (self.schema_config and self.schema_config.value_schema_type):
             deserializer = get_deserializer(
-                self.deserialize.value_schema_type,
-                self.deserialize.value_schema_str,
+                self.schema_config.value_schema_type,
+                self.schema_config.value_schema_str,
             )
             deserialized_value = deserializer.deserialize(value)
 
-            if self.deserialize.value_output_serializer:
+            if self.schema_config.value_output_serializer:
                 return serialize_to_output_type(
                     deserialized_value,
-                    self.deserialize.value_output_serializer,
+                    self.schema_config.value_output_serializer,
                 )
 
             return deserialized_value
@@ -80,35 +81,35 @@ def original_key(self) -> str | None:
         return self.get("key")
 
     @property
-    def headers(self) -> list[dict[str, list[int]]]:
+    def original_headers(self) -> list[dict[str, list[int]]]:
         """The raw Kafka record headers."""
-        return CaseInsensitiveDict((k, bytes(v)) for chunk in self.headers for k, v in chunk.items())
+        return self["headers"]
 
-    @property
-    def original_headers(self) -> dict[str, bytes]:
+    @cached_property
+    def headers(self) -> dict[str, bytes]:
         """Decodes the headers as a single dictionary."""
-        return self["headers"]
+        return CaseInsensitiveDict((k, bytes(v)) for chunk in self.original_headers for k, v in chunk.items())
 
 
-class ConsumerRecords(KafkaEvent):
+class ConsumerRecords(KafkaEventBase):
     """Self-managed or MSK Apache Kafka event trigger
     Documentation:
     --------------
     - https://docs.aws.amazon.com/lambda/latest/dg/with-kafka.html
     - https://docs.aws.amazon.com/lambda/latest/dg/with-msk.html
     """
 
-    def __init__(self, data: dict[str, Any], deserialize: SchemaConfig | None = None):
+    def __init__(self, data: dict[str, Any], schema_config: SchemaConfig | None = None):
         super().__init__(data)
         self._records: Iterator[ConsumerRecordRecords] | None = None
-        self.deserialize = deserialize
+        self.schema_config = schema_config
 
     @property
     def records(self) -> Iterator[ConsumerRecordRecords]:
         """The Kafka records."""
         for chunk in self["records"].values():
             for record in chunk:
-                yield ConsumerRecordRecords(data=record, deserialize=self.deserialize)
+                yield ConsumerRecordRecords(data=record, schema_config=self.schema_config)
 
     @property
     def record(self) -> ConsumerRecordRecords:

aws_lambda_powertools/utilities/kafka_consumer/deserializer/avro.py

@@ -30,7 +30,7 @@ def __init__(self, schema_str: str):
                 f"Invalid Avro schema. Please ensure the provided avro schema is valid: {type(e).__name__}: {str(e)}",
             ) from e
 
-    def deserialize(self, data: bytes | str) -> dict[str, Any]:
+    def deserialize(self, data: bytes | str) -> object:
         """
         Deserialize Avro binary data to a Python dictionary.
 

aws_lambda_powertools/utilities/kafka_consumer/deserializer/base.py

@@ -2,7 +2,7 @@
 
 import base64
 from abc import ABC, abstractmethod
-from typing import Any
+from typing import Any, overload
 
 
 class DeserializerBase(ABC):
@@ -29,7 +29,7 @@ class DeserializerBase(ABC):
     """
 
     @abstractmethod
-    def deserialize(self, data: bytes | str) -> dict[str, Any]:
+    def deserialize(self, data: bytes | str) -> dict[str, Any] | str | object:
         """
         Deserialize input data to a Python dictionary.
 

aws_lambda_powertools/utilities/kafka_consumer/deserializer/default.py

@@ -14,7 +14,7 @@ class DefaultDeserializer(DeserializerBase):
     no customized deserialization is needed or when handling raw data formats.
     """
 
-    def deserialize(self, data: bytes | str) -> dict[str, Any]:
+    def deserialize(self, data: bytes | str) -> str:
         """
         Return the input data base64 decoded.
 

aws_lambda_powertools/utilities/kafka_consumer/deserializer/json.py

@@ -46,8 +46,7 @@ def deserialize(self, data: bytes | str) -> dict:
         ...     print(f"Failed to deserialize: {e}")
         """
         try:
-            value = self._decode_input(data)
-            return json.loads(base64.b64decode(value).decode("utf-8"))
+            return json.loads(base64.b64decode(data).decode("utf-8"))
         except Exception as e:
             raise KafkaConsumerDeserializationError(
                 f"JSON deserialization error: {type(e).__name__}: {str(e)}",

aws_lambda_powertools/utilities/kafka_consumer/exceptions.py

@@ -14,3 +14,9 @@ class KafkaConsumerMissingSchemaError(Exception):
     """
     Error raised when a required schema is not provided.
     """
+
+
+class KafkaConsumerOutputSerializerError(Exception):
+    """
+    Error raised when output serializer fails.
+    """

aws_lambda_powertools/utilities/kafka_consumer/schema_config.py

@@ -57,10 +57,10 @@ class SchemaConfig:
 
     def __init__(
         self,
-        value_schema_type: Literal["AVRO", "PROTOBUF", "JSON"] = None,
+        value_schema_type: Literal["AVRO", "PROTOBUF", "JSON"] | None = None,
         value_schema: str | None = None,
         value_output_serializer: Any | None = None,
-        key_schema_type: Literal["AVRO", "PROTOBUF", "JSON", None] = None,
+        key_schema_type: Literal["AVRO", "PROTOBUF", "JSON", None] | None = None,
         key_schema: str | None = None,
         key_output_serializer: Any | None = None,
     ):

aws_lambda_powertools/utilities/kafka_consumer/serialization/base.py

@@ -51,4 +51,4 @@ def serialize(self, data: dict[str, Any], output_class: type[T] | None = None) -
             An instance of output_class if provided, otherwise a processed dictionary.
             The generic type T represents the type of the output_class.
         """
-        pass
+        raise NotImplementedError("Subclasses must implement this method")

aws_lambda_powertools/utilities/kafka_consumer/serialization/custom_dict.py

@@ -2,6 +2,7 @@
 
 from typing import TYPE_CHECKING, Any
 
+from aws_lambda_powertools.utilities.kafka_consumer.exceptions import KafkaConsumerOutputSerializerError
 from aws_lambda_powertools.utilities.kafka_consumer.serialization.base import OutputSerializerBase
 
 if TYPE_CHECKING:
@@ -13,14 +14,11 @@ def serialize(self, data: dict[str, Any], output_class: type[T] | None = None) -
         if output_class is None:
             return data
 
-        if not hasattr(output_class, "to_dict") and not hasattr(output_class, "from_dict"):
-            raise ValueError("Output class must have to_dict or from_dict method")
-
-        if hasattr(output_class, "from_dict"):
-            return output_class.from_dict(data)
+        if not hasattr(output_class, "to_dict"):
+            raise KafkaConsumerOutputSerializerError("The output serialization class must have to_dict method")
 
         # Instantiate and then populate
-        instance = output_class()
+        instance = output_class
         for key, value in data.items():
             setattr(instance, key, value)
         return instance