Adding docstring

Author: leandrodamascena

aws_lambda_powertools/utilities/kafka_consumer/consumer_records.py

@@ -27,8 +27,8 @@ def key(self) -> Any:
         key = self.get("key")
         if key and (self.deserialize and self.deserialize.key_schema_type):
             deserializer = get_deserializer(
-                self.deserialize.value_schema_type,
-                self.deserialize.value_schema_str,
+                self.deserialize.key_schema_type,
+                self.deserialize.key_schema_str,
             )
             deserialized_key = deserializer.deserialize(key)
 

aws_lambda_powertools/utilities/kafka_consumer/deserializer/avro.py

@@ -8,25 +8,65 @@
 
 from aws_lambda_powertools.utilities.kafka_consumer.deserializer.base import DeserializerBase
 from aws_lambda_powertools.utilities.kafka_consumer.exceptions import (
-    KafkaConsumerAvroMissingSchemaError,
+    KafkaConsumerAvroSchemaParserError,
     KafkaConsumerDeserializationError,
 )
 
 
 class AvroDeserializer(DeserializerBase):
+    """
+    Deserializer for Apache Avro formatted data.
+
+    This class provides functionality to deserialize Avro binary data using
+    a provided Avro schema definition.
+    """
+
     def __init__(self, schema_str: str):
-        if not schema_str:
-            raise KafkaConsumerAvroMissingSchemaError("Schema string must be provided for Avro deserialization")
-        self.parsed_schema = parse_schema(schema_str)
-        self.reader = DatumReader(self.parsed_schema)
+        try:
+            self.parsed_schema = parse_schema(schema_str)
+            self.reader = DatumReader(self.parsed_schema)
+        except Exception as e:
+            raise KafkaConsumerAvroSchemaParserError(
+                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]:
+        """
+        Deserialize Avro binary data to a Python dictionary.
+
+        Parameters
+        ----------
+        data : bytes or str
+            The Avro binary data to deserialize. If provided as a string,
+            it will be decoded to bytes first.
+
+        Returns
+        -------
+        dict[str, Any]
+            Deserialized data as a dictionary.
+
+        Raises
+        ------
+        KafkaConsumerDeserializationError
+            When the data cannot be deserialized according to the schema,
+            typically due to data format incompatibility.
+
+        Examples
+        --------
+        >>> deserializer = AvroDeserializer(schema_str)
+        >>> avro_data = b'...'  # binary Avro data
+        >>> try:
+        ...     result = deserializer.deserialize(avro_data)
+        ...     # Process the deserialized data
+        ... except KafkaConsumerDeserializationError as e:
+        ...     print(f"Failed to deserialize: {e}")
+        """
         try:
             value = self._decode_input(data)
             bytes_reader = io.BytesIO(value)
             decoder = BinaryDecoder(bytes_reader)
             return self.reader.read(decoder)
-        except (TypeError, ValueError) as e:
+        except Exception as e:
             raise KafkaConsumerDeserializationError(
-                f"Avro deserialization error: {type(e).__name__}: {str(e)}",
+                f"Error trying to deserializer avro data - {type(e).__name__}: {str(e)}",
             ) from e

aws_lambda_powertools/utilities/kafka_consumer/deserializer/base.py

@@ -6,9 +6,47 @@
 
 
 class DeserializerBase(ABC):
+    """
+    Abstract base class for deserializers.
+
+    This class defines the interface for all deserializers in the Kafka consumer utility
+    and provides a common method for decoding input data.
+
+    Methods
+    -------
+    deserialize(data)
+        Abstract method that must be implemented by subclasses to deserialize data.
+    _decode_input(data)
+        Helper method to decode input data to bytes.
+
+    Examples
+    --------
+    >>> class MyDeserializer(DeserializerBase):
+    ...     def deserialize(self, data: bytes | str) -> dict[str, Any]:
+    ...         value = self._decode_input(data)
+    ...         # Custom deserialization logic here
+    ...         return {"key": "value"}
+    """
+
     @abstractmethod
     def deserialize(self, data: bytes | str) -> dict[str, Any]:
-        pass
+        """
+        Deserialize input data to a Python dictionary.
+
+        This abstract method must be implemented by subclasses to provide
+        specific deserialization logic.
+
+        Parameters
+        ----------
+        data : bytes or str
+            The data to deserialize, either as bytes or as a string.
+
+        Returns
+        -------
+        dict[str, Any]
+            The deserialized data as a dictionary.
+        """
+        raise NotImplementedError("Subclasses must implement the deserialize method")
 
     def _decode_input(self, data: bytes | str) -> bytes:
         if isinstance(data, str):

aws_lambda_powertools/utilities/kafka_consumer/deserializer/deserializer.py

@@ -2,11 +2,59 @@
 
 from typing import TYPE_CHECKING, Any
 
+from aws_lambda_powertools.utilities.kafka_consumer.deserializer.json import JsonDeserializer
+from aws_lambda_powertools.utilities.kafka_consumer.deserializer.no_op import NoOpDeserializer
+
 if TYPE_CHECKING:
     from aws_lambda_powertools.utilities.kafka_consumer.deserializer.base import DeserializerBase
 
 
-def get_deserializer(schema_type: str, schema_value: Any) -> DeserializerBase:
+def get_deserializer(schema_type: str | object, schema_value: Any) -> DeserializerBase:
+    """
+    Factory function to get the appropriate deserializer based on schema type.
+
+    This function creates and returns a deserializer instance that corresponds to the
+    specified schema type. It handles lazy imports for optional dependencies.
+
+    Parameters
+    ----------
+    schema_type : str
+        The type of schema to use for deserialization.
+        Supported values are: "AVRO", "PROTOBUF", "JSON", or any other value for no-op.
+    schema_value : Any
+        The schema definition to use for deserialization. The format depends on the
+        schema_type:
+        - For "AVRO": A string containing the Avro schema definition
+        - For "PROTOBUF": A object containing the Protobuf schema definition
+        - For "JSON": Not used (can be None)
+        - For other types: Not used (can be None)
+
+    Returns
+    -------
+    DeserializerBase
+        An instance of a deserializer that implements the DeserializerBase interface.
+
+    Examples
+    --------
+    >>> # Get an Avro deserializer
+    >>> avro_schema = '''
+    ...     {
+    ...       "type": "record",
+    ...       "name": "User",
+    ...       "fields": [
+    ...         {"name": "name", "type": "string"},
+    ...         {"name": "age", "type": "int"}
+    ...       ]
+    ...     }
+    ... '''
+    >>> deserializer = get_deserializer("AVRO", avro_schema)
+    >>>
+    >>> # Get a JSON deserializer
+    >>> json_deserializer = get_deserializer("JSON", None)
+    >>>
+    >>> # Get a no-op deserializer for raw data
+    >>> no_op_deserializer = get_deserializer("RAW", None)
+    """
     if schema_type == "AVRO":
         # Import here to avoid dependency if not used
         from aws_lambda_powertools.utilities.kafka_consumer.deserializer.avro import AvroDeserializer
@@ -18,9 +66,6 @@ def get_deserializer(schema_type: str, schema_value: Any) -> DeserializerBase:
 
         return ProtobufDeserializer(schema_value)
     elif schema_type == "JSON":
-        # Import here to avoid dependency if not used
-        from aws_lambda_powertools.utilities.kafka_consumer.deserializer.json import JsonDeserializer
-
         return JsonDeserializer()
-    else:
-        raise ValueError(f"Invalid schema_type: {schema_type}")
+
+    return NoOpDeserializer()

aws_lambda_powertools/utilities/kafka_consumer/deserializer/json.py

@@ -7,7 +7,43 @@
 
 
 class JsonDeserializer(DeserializerBase):
+    """
+    Deserializer for JSON formatted data.
+
+    This class provides functionality to deserialize JSON data from bytes or string
+    into Python dictionaries.
+    """
+
     def deserialize(self, data: bytes | str) -> dict:
+        """
+        Deserialize JSON data to a Python dictionary.
+
+        Parameters
+        ----------
+        data : bytes or str
+            The JSON data to deserialize. If provided as bytes, it will be decoded as UTF-8.
+            If provided as a string, it's assumed to be base64-encoded and will be decoded first.
+
+        Returns
+        -------
+        dict
+            Deserialized data as a dictionary.
+
+        Raises
+        ------
+        KafkaConsumerDeserializationError
+            When the data cannot be deserialized as valid JSON.
+
+        Examples
+        --------
+        >>> deserializer = JsonDeserializer()
+        >>> json_data = '{"key": "value", "number": 123}'
+        >>> try:
+        ...     result = deserializer.deserialize(json_data)
+        ...     print(result["key"])  # Output: value
+        ... except KafkaConsumerDeserializationError as e:
+        ...     print(f"Failed to deserialize: {e}")
+        """
         try:
             value = self._decode_input(data)
             return json.loads(value.decode("utf-8"))

aws_lambda_powertools/utilities/kafka_consumer/deserializer/no_op.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from typing import Any
+
+from aws_lambda_powertools.utilities.kafka_consumer.deserializer.base import DeserializerBase
+
+
+class NoOpDeserializer(DeserializerBase):
+    """
+    A pass-through deserializer that performs no transformation on the input data.
+
+    This deserializer simply returns the input data unchanged, which is useful when
+    no deserialization is needed or when handling raw data formats.
+    """
+
+    def deserialize(self, data: bytes | str) -> dict[str, Any]:
+        """
+        Return the input data unchanged.
+
+        This method implements the deserialize interface but performs no transformation,
+        simply returning the input data as-is.
+
+        Parameters
+        ----------
+        data : bytes or str
+            The input data to "deserialize".
+
+        Returns
+        -------
+        dict[str, Any]
+            The input data unchanged. Note that despite the type annotation,
+            this method returns the exact same object that was passed in,
+            preserving its original type.
+
+        Example
+        --------
+        >>> deserializer = NoOpDeserializer()
+        >>>
+        >>> # With string input
+        >>> string_data = "Hello, world!"
+        >>> result = deserializer.deserialize(string_data)
+        >>> print(result == string_data)  # Output: True
+        >>>
+        >>> # With bytes input
+        >>> bytes_data = b"Binary data"
+        >>> result = deserializer.deserialize(bytes_data)
+        >>> print(result == bytes_data)  # Output: True
+        """
+        return data

aws_lambda_powertools/utilities/kafka_consumer/deserializer/protobuf.py

@@ -11,10 +11,49 @@
 
 
 class ProtobufDeserializer(DeserializerBase):
+    """
+    Deserializer for Protocol Buffer formatted data.
+
+    This class provides functionality to deserialize Protocol Buffer binary data
+    into Python dictionaries using the provided Protocol Buffer message class.
+    """
+
     def __init__(self, message_class: Any):
         self.message_class = message_class
 
     def deserialize(self, data: bytes | str) -> dict:
+        """
+        Deserialize Protocol Buffer binary data to a Python dictionary.
+
+        Parameters
+        ----------
+        data : bytes or str
+            The Protocol Buffer binary data to deserialize. If provided as a string,
+            it's assumed to be base64-encoded and will be decoded first.
+
+        Returns
+        -------
+        dict
+            Deserialized data as a dictionary with field names preserved from the
+            Protocol Buffer definition.
+
+        Raises
+        ------
+        KafkaConsumerDeserializationError
+            When the data cannot be deserialized according to the message class,
+            typically due to data format incompatibility or incorrect message class.
+
+        Example
+        --------
+        >>> # Assuming proper protobuf setup
+        >>> deserializer = ProtobufDeserializer(my_proto_module.MyMessage)
+        >>> proto_data = b'...'  # binary protobuf data
+        >>> try:
+        ...     result = deserializer.deserialize(proto_data)
+        ...     # Process the deserialized dictionary
+        ... except KafkaConsumerDeserializationError as e:
+        ...     print(f"Failed to deserialize: {e}")
+        """
         try:
             value = self._decode_input(data)
             message = self.message_class()

aws_lambda_powertools/utilities/kafka_consumer/exceptions.py

@@ -1,16 +1,16 @@
-class KafkaConsumerAvroSchemaMismatchError(Exception):
+class KafkaConsumerAvroSchemaParserError(Exception):
     """
-    Avro schema mismatch
+    Error raised when parsing Avro schema definition fails.
     """
 
 
 class KafkaConsumerDeserializationError(Exception):
     """
-    Avro schema impossible to deserialize
+    Error raised when message deserialization fails.
     """
 
 
-class KafkaConsumerAvroMissingSchemaError(Exception):
+class KafkaConsumerMissingSchemaError(Exception):
     """
-    Avro schema mismatch
+    Error raised when a required schema is not provided.
     """