Added some docstrings and typing

Author: seshubaws

aws_lambda_powertools/utilities/data_masking/base.py

@@ -1,11 +1,11 @@
 import json
-from typing import Union
+from typing import Union, Optional
 
 from aws_lambda_powertools.utilities.data_masking.provider import BaseProvider
 
 
 class DataMasking:
-    def __init__(self, provider=None):
+    def __init__(self, provider: Optional[BaseProvider] = None):
         self.provider = provider or BaseProvider()
 
     def encrypt(self, data, fields=None, **provider_options):
@@ -18,37 +18,66 @@ def mask(self, data, fields=None, **provider_options) -> str:
         return self._apply_action(data, fields, self.provider.mask, **provider_options)
 
     def _apply_action(self, data, fields, action, **provider_options):
+        """
+        Helper method to determine whether to apply a given action to the entire input data
+        or to specific fields if the 'fields' argument is specified.
+
+        Parameters
+        ----------
+        data : any
+            The input data to process.
+        fields : Optional[List[any]] = None
+            A list of fields to apply the action to. If 'None', the action is applied to the entire 'data'.
+        action : Callable
+           The action to apply to the data. It should be a callable that performs an operation on the data
+           and returns the modified value.
+
+        Returns
+        -------
+        any
+            The modified data after applying the action.
+        """
+
         if fields is not None:
             return self._apply_action_to_fields(data, fields, action, **provider_options)
         else:
             return action(data, **provider_options)
 
     def _apply_action_to_fields(self, data: Union[dict, str], fields, action, **provider_options) -> Union[dict, str]:
         """
-        Apply the specified action to the specified fields in the input data.
-
-        This method is takes the input data, which can be either a dictionary or a JSON string representation
+        This method takes the input data, which can be either a dictionary or a JSON string representation
         of a dictionary, and applies a mask, an encryption, or a decryption to the specified fields.
 
-        Parameters:
-            data (Union[dict, str]): The input data to process. It can be either a dictionary or a JSON string
+        Parameters
+        ----------
+            data : Union[dict, str])
+                The input data to process. It can be either a dictionary or a JSON string
                 representation of a dictionary.
-            fields (list): A list of fields to apply the action to. Each field can be specified as a string or
+            fields : List
+                A list of fields to apply the action to. Each field can be specified as a string or
                 a list of strings representing nested keys in the dictionary.
-            action (callable): The action to apply to the fields. It should be a callable that takes the current
+            action : Callable
+                The action to apply to the fields. It should be a callable that takes the current
                 value of the field as the first argument and any additional arguments that might be required
                 for the action. It performs an operation on the current value using the provided arguments and
                 returns the modified value.
-            **provider_options: Additional keyword arguments to pass to the 'action' function.
+            **provider_options: 
+                Additional keyword arguments to pass to the 'action' function.
 
-        Returns:
-            str: A JSON string representation of the modified dictionary after applying the action to the
+        Returns
+        -------
+            str
+                A JSON string representation of the modified dictionary after applying the action to the
             specified fields.
 
-        Raises:
-            ValueError: If 'fields' parameter is None.
-            TypeError: If the 'data' parameter is not a dictionary or a JSON string representation of a dictionary.
-            KeyError: If specified 'fields' do not exist in input data
+        Raises
+        -------
+            ValueError
+                If 'fields' parameter is None.
+            TypeError
+                If the 'data' parameter is not a dictionary or a JSON string representation of a dictionary.
+            KeyError
+                If specified 'fields' do not exist in input data
         """
 
         if fields is None:

aws_lambda_powertools/utilities/data_masking/provider.py

@@ -1,22 +1,20 @@
-from abc import abstractmethod
+from abc import ABC
 from collections.abc import Iterable
 from typing import Union
 
 from aws_lambda_powertools.utilities.data_masking.constants import DATA_MASKING_STRING
 
 
-class BaseProvider:
+class BaseProvider(ABC):
     """
     When you try to create an instance of a subclass that does not implement the encrypt method,
     you will get a NotImplementedError with a message that says the method is not implemented:
     """
 
-    @abstractmethod
     def encrypt(self, data) -> Union[bytes, str]:
         raise NotImplementedError("Subclasses must implement encrypt()")
 
-    @abstractmethod
-    def decrypt(self, data) -> Union[bytes, str]:
+    def decrypt(self, data) -> any:
         raise NotImplementedError("Subclasses must implement decrypt()")
 
     def mask(self, data) -> str:

aws_lambda_powertools/utilities/data_masking/providers/aws_encryption_sdk.py

@@ -39,7 +39,9 @@ def __call__(cls, *args, **provider_options):
 class AwsEncryptionSdkProvider(BaseProvider):
     """
     The AwsEncryptionSdkProvider is to be used as a Provider for the Datamasking class.
-    Example:
+    
+    Example
+    -------
         >>> data_masker = DataMasking(provider=AwsEncryptionSdkProvider(keys=[keyARN1, keyARN2,...,]))
         >>> encrypted_data = data_masker.encrypt("a string")
         "encrptedBase64String"
@@ -73,14 +75,16 @@ def encrypt(self, data: Union[bytes, str], **provider_options) -> str:
         """
         Encrypt data using the AwsEncryptionSdkProvider.
 
-        Parameters:
-            - data (Union[bytes, str]):
+        Parameters
+        -------
+            data : Union[bytes, str]
                 The data to be encrypted.
-            - provider_options:
+            provider_options
                 Additional options for the aws_encryption_sdk.EncryptionSDKClient
 
-        Returns:
-            - ciphertext (str):
+        Returns
+        -------
+            ciphertext : str
                 The encrypted data, as a base64-encoded string.
         """
         ciphertext, _ = self.client.encrypt(source=data, materials_manager=self.cache_cmm, **provider_options)
@@ -91,14 +95,16 @@ def decrypt(self, data: str, **provider_options) -> bytes:
         """
         Decrypt data using AwsEncryptionSdkProvider.
 
-        Parameters:
-            - data (Union[bytes, str]):
-                The encrypted data, as a base64-encoded string.
-            - provider_options:
+        Parameters
+        -------
+            data : Union[bytes, str]
+                The encrypted data, as a base64-encoded string
+            provider_options
                 Additional options for the aws_encryption_sdk.EncryptionSDKClient
 
-        Returns:
-            - ciphertext (bytes):
+        Returns
+        -------
+            ciphertext : bytes
                 The decrypted data in bytes
         """
         ciphertext_decoded = base64.b64decode(data)