Fixing typing

aws_lambda_powertools/utilities/_data_masking/base.py

@@ -53,7 +53,7 @@ def lambda_handler(event, context):
     def __init__(self, provider: Optional[BaseProvider] = None):
         self.provider = provider or BaseProvider()
 
-    def encrypt(self, data, fields=None, **provider_options) -> str:
+    def encrypt(self, data, fields=None, **provider_options) -> str | dict:
         return self._apply_action(data, fields, self.provider.encrypt, **provider_options)
 
     def decrypt(self, data, fields=None, **provider_options) -> Any:

aws_lambda_powertools/utilities/_data_masking/provider/base.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import json
 from typing import Any, Iterable, Union
 
@@ -75,7 +77,7 @@ def default_json_serializer(self, data):
     def default_json_deserializer(self, data):
         return json.loads(data.decode("utf-8"))
 
-    def encrypt(self, data) -> str:
+    def encrypt(self, data) -> str | dict:
         raise NotImplementedError("Subclasses must implement encrypt()")
 
     def decrypt(self, data) -> Any:

examples/data_masking/src/getting_started_decrypt_context.py

@@ -21,7 +21,7 @@ def lambda_handler(event: dict, context: LambdaContext) -> dict:
 
     logger.info("Decrypting email field")
 
-    decrypted: dict = data_masker.encrypt(
+    decrypted = data_masker.decrypt(
         data,
         fields=["email"],
         tenant_id=event.get("tenant_id", ""),  # (1)!

examples/data_masking/src/getting_started_encrypt_context.py

@@ -17,12 +17,12 @@
 
 
 @logger.inject_lambda_context
-def lambda_handler(event: dict, context: LambdaContext) -> dict:
+def lambda_handler(event: dict, context: LambdaContext) -> dict | str:
     data = event.get("body", {})
 
     logger.info("Encrypting email field")
 
-    encrypted: dict = data_masker.encrypt(
+    encrypted: dict | str = data_masker.encrypt(
         data,
         fields=["email"],
         data_classification="confidential",

examples/data_masking/src/getting_started_encrypt_data.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+from typing import Iterable
 
 from aws_lambda_powertools import Logger
 from aws_lambda_powertools.utilities._data_masking import DataMasking
@@ -16,11 +17,11 @@
 
 
 @logger.inject_lambda_context
-def lambda_handler(event: dict, context: LambdaContext) -> dict:
+def lambda_handler(event: dict, context: LambdaContext) -> Iterable | str:
     data = event.get("body", {})
 
     logger.info("Encrypting fields email, address.street, and company_address")
 
-    encrypted: dict = data_masker.encrypt(data, fields=["email", "address.street", "company_address"])  # (2)!
+    encrypted: Iterable = data_masker.encrypt(data, fields=["email", "address.street", "company_address"])  # (2)!
 
     return encrypted

examples/data_masking/src/getting_started_mask_data.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import Iterable
+
 from aws_lambda_powertools import Logger
 from aws_lambda_powertools.utilities._data_masking import DataMasking
 from aws_lambda_powertools.utilities.typing import LambdaContext
@@ -9,11 +11,11 @@
 
 
 @logger.inject_lambda_context
-def lambda_handler(event: dict, context: LambdaContext) -> dict:
+def lambda_handler(event: dict, context: LambdaContext) -> Iterable | str:
     data = event.get("body")
 
     logger.info("Masking fields email, address.street, and company_address")
 
-    masked: dict = data_masker.mask(data, fields=["email", "address.street", "company_address"])  # (1)!
+    masked: Iterable = data_masker.mask(data, fields=["email", "address.street", "company_address"])  # (1)!
 
     return masked