From 0cb967a3563a722a9c220a2a4c20851d9df934d6 Mon Sep 17 00:00:00 2001 From: Seshu Brahma Date: Tue, 12 Dec 2023 16:47:35 -0800 Subject: [PATCH] Add return types for mask() --- aws_lambda_powertools/utilities/_data_masking/base.py | 4 ++-- .../utilities/_data_masking/provider/base.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/aws_lambda_powertools/utilities/_data_masking/base.py b/aws_lambda_powertools/utilities/_data_masking/base.py index 0b395fdbb07..fc6f827458d 100644 --- 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 index a293c6aff9a..f1488fe542d 100644 --- 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)):