From 483c1b4425ff6e034299c06dd50bb948b5c966ea Mon Sep 17 00:00:00 2001
From: Leandro Damascena <lcdama@amazon.pt>
Date: Mon, 18 Dec 2023 20:00:00 +0000
Subject: [PATCH] Fixing typing

---
 aws_lambda_powertools/utilities/_data_masking/base.py       | 2 +-
 .../utilities/_data_masking/provider/base.py                | 4 +++-
 .../data_masking/src/getting_started_decrypt_context.py     | 2 +-
 .../data_masking/src/getting_started_encrypt_context.py     | 4 ++--
 examples/data_masking/src/getting_started_encrypt_data.py   | 5 +++--
 examples/data_masking/src/getting_started_mask_data.py      | 6 ++++--
 6 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/aws_lambda_powertools/utilities/_data_masking/base.py b/aws_lambda_powertools/utilities/_data_masking/base.py
index 96905c993aa..08018e4f103 100644
--- a/aws_lambda_powertools/utilities/_data_masking/base.py
+++ b/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:
diff --git a/aws_lambda_powertools/utilities/_data_masking/provider/base.py b/aws_lambda_powertools/utilities/_data_masking/provider/base.py
index ad16216e8b5..55d08a9be32 100644
--- a/aws_lambda_powertools/utilities/_data_masking/provider/base.py
+++ b/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:
diff --git a/examples/data_masking/src/getting_started_decrypt_context.py b/examples/data_masking/src/getting_started_decrypt_context.py
index fd46309cec5..56c515e718b 100644
--- a/examples/data_masking/src/getting_started_decrypt_context.py
+++ b/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)!
diff --git a/examples/data_masking/src/getting_started_encrypt_context.py b/examples/data_masking/src/getting_started_encrypt_context.py
index d34c8fd7d12..1442f8d5805 100644
--- a/examples/data_masking/src/getting_started_encrypt_context.py
+++ b/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",
diff --git a/examples/data_masking/src/getting_started_encrypt_data.py b/examples/data_masking/src/getting_started_encrypt_data.py
index 425e634dac0..2246662e4e4 100644
--- a/examples/data_masking/src/getting_started_encrypt_data.py
+++ b/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
diff --git a/examples/data_masking/src/getting_started_mask_data.py b/examples/data_masking/src/getting_started_mask_data.py
index b31b1e044a8..84e6dde59ca 100644
--- a/examples/data_masking/src/getting_started_mask_data.py
+++ b/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