refactor(data_masking): add from __future__ import annotations

aws_lambda_powertools/utilities/data_masking/base.py

@@ -3,8 +3,7 @@
 import functools
 import logging
 import warnings
-from numbers import Number
-from typing import Any, Callable, Mapping, Optional, Sequence, Union, overload
+from typing import TYPE_CHECKING, Any, Callable, Mapping, Sequence, overload
 
 from jsonpath_ng.ext import parse
 
@@ -14,6 +13,9 @@
 )
 from aws_lambda_powertools.utilities.data_masking.provider import BaseProvider
 
+if TYPE_CHECKING:
+    from numbers import Number
+
 logger = logging.getLogger(__name__)
 
 
@@ -43,7 +45,7 @@ def lambda_handler(event, context):
 
     def __init__(
         self,
-        provider: Optional[BaseProvider] = None,
+        provider: BaseProvider | None = None,
         raise_on_missing_field: bool = True,
     ):
         self.provider = provider or BaseProvider()
@@ -111,7 +113,7 @@ def _apply_action(
         ----------
         data : str | dict
             The input data to process.
-        fields : Optional[List[str]]
+        fields : list[str] | 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
@@ -142,21 +144,21 @@ def _apply_action(
 
     def _apply_action_to_fields(
         self,
-        data: Union[dict, str],
+        data: dict | str,
         fields: list,
         action: Callable,
         provider_options: dict | None = None,
         **encryption_context: str,
-    ) -> Union[dict, str]:
+    ) -> dict | str:
         """
         This method takes the input data, which can be either a dictionary or a JSON string,
         and erases, encrypts, or decrypts the specified fields.
 
         Parameters
         ----------
-            data : Union[dict, str])
+            data : dict | str)
                 The input data to process. It can be either a dictionary or a JSON string.
-            fields : List
+            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

aws_lambda_powertools/utilities/data_masking/provider/base.py

@@ -24,7 +24,7 @@ def encrypt(self, data) -> str:
         def decrypt(self, data) -> Any:
             # Implementation logic for data decryption
 
-        def erase(self, data) -> Union[str, Iterable]:
+        def erase(self, data) -> str | Iterable:
             # Implementation logic for data masking
             pass
 

aws_lambda_powertools/utilities/data_masking/provider/kms/aws_encryption_sdk.py

@@ -4,7 +4,7 @@
 import json
 import logging
 from binascii import Error
-from typing import Any, Callable, List
+from typing import Any, Callable
 
 import botocore
 from aws_encryption_sdk import (
@@ -18,7 +18,6 @@
     GenerateKeyError,
     NotSupportedError,
 )
-from aws_encryption_sdk.structures import MessageHeader
 
 from aws_lambda_powertools.shared.functions import (
     base64_decode,
@@ -76,7 +75,7 @@ def lambda_handler(event, context):
 
     def __init__(
         self,
-        keys: List[str],
+        keys: list[str],
         key_provider=None,
         local_cache_capacity: int = CACHE_CAPACITY,
         max_cache_age_seconds: float = MAX_CACHE_AGE_SECONDS,
@@ -112,7 +111,7 @@ class KMSKeyProvider:
 
     def __init__(
         self,
-        keys: List[str],
+        keys: list[str],
         json_serializer: Callable[..., str],
         json_deserializer: Callable[[str], Any],
         local_cache_capacity: int = CACHE_CAPACITY,
@@ -143,7 +142,7 @@ def encrypt(self, data: Any, provider_options: dict | None = None, **encryption_
 
         Parameters
         -------
-            data : Union[bytes, str]
+            data : Any
                 The data to be encrypted.
             provider_options : dict
                 Additional options for the aws_encryption_sdk.EncryptionSDKClient
@@ -180,7 +179,7 @@ def decrypt(self, data: str, provider_options: dict | None = None, **encryption_
 
         Parameters
         -------
-            data : Union[bytes, str]
+            data : str
                 The encrypted data, as a base64-encoded string
             provider_options
                 Additional options for the aws_encryption_sdk.EncryptionSDKClient
@@ -201,8 +200,6 @@ def decrypt(self, data: str, provider_options: dict | None = None, **encryption_
             )
 
         try:
-            decryptor_header: MessageHeader
-
             ciphertext, decryptor_header = self.client.decrypt(
                 source=ciphertext_decoded,
                 key_provider=self.key_provider,