Add return types for mask()

seshubaws · Dec 13, 2023 · 0cb967a · 0cb967a
1 parent ecf505e
commit 0cb967a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/aws_lambda_powertools/utilities/_data_masking/base.py b/aws_lambda_powertools/utilities/_data_masking/base.py
@@ -1,5 +1,5 @@
 import json
-from typing import Callable, Optional, Union
+from typing import Callable, Iterable, Optional, Union
 
 from aws_lambda_powertools.utilities._data_masking.exceptions import DataMaskingUnsupportedTypeError
 from aws_lambda_powertools.utilities._data_masking.provider import BaseProvider
@@ -51,7 +51,7 @@ def encrypt(self, data, fields=None, **provider_options):
     def decrypt(self, data, fields=None, **provider_options):
         return self._apply_action(data, fields, self.provider.decrypt, **provider_options)
 
-    def mask(self, data, fields=None, **provider_options):
+    def mask(self, data, fields=None, **provider_options) -> Union[str, Iterable]:
         return self._apply_action(data, fields, self.provider.mask, **provider_options)
 
     def _apply_action(self, data, fields, action: Callable, **provider_options):

diff --git a/aws_lambda_powertools/utilities/_data_masking/provider/base.py b/aws_lambda_powertools/utilities/_data_masking/provider/base.py
@@ -1,5 +1,5 @@
 import json
-from typing import Any
+from typing import Any, Iterable, Union
 
 from aws_lambda_powertools.utilities._data_masking.constants import DATA_MASKING_STRING
 
@@ -26,7 +26,7 @@ def encrypt(self, data) -> str:
     def decrypt(self, data) -> Any:
         raise NotImplementedError("Subclasses must implement decrypt()")
 
-    def mask(self, data) -> Any:
+    def mask(self, data) -> Union[str, Iterable]:
         if isinstance(data, (str, dict, bytes)):
             return DATA_MASKING_STRING
         elif isinstance(data, (list, tuple, set)):