Skip to content

Commit 456bf82

Browse files
authored
refactor(data_masking): add from __future__ import annotations (#4945)
and update code according to ruff rules TCH, UP006, UP007, UP037 and FA100.
1 parent 56914af commit 456bf82

File tree

3 files changed

+16
-17
lines changed

3 files changed

+16
-17
lines changed

aws_lambda_powertools/utilities/data_masking/base.py

+10-8
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,7 @@
33
import functools
44
import logging
55
import warnings
6-
from numbers import Number
7-
from typing import Any, Callable, Mapping, Optional, Sequence, Union, overload
6+
from typing import TYPE_CHECKING, Any, Callable, Mapping, Sequence, overload
87

98
from jsonpath_ng.ext import parse
109

@@ -14,6 +13,9 @@
1413
)
1514
from aws_lambda_powertools.utilities.data_masking.provider import BaseProvider
1615

16+
if TYPE_CHECKING:
17+
from numbers import Number
18+
1719
logger = logging.getLogger(__name__)
1820

1921

@@ -43,7 +45,7 @@ def lambda_handler(event, context):
4345

4446
def __init__(
4547
self,
46-
provider: Optional[BaseProvider] = None,
48+
provider: BaseProvider | None = None,
4749
raise_on_missing_field: bool = True,
4850
):
4951
self.provider = provider or BaseProvider()
@@ -111,7 +113,7 @@ def _apply_action(
111113
----------
112114
data : str | dict
113115
The input data to process.
114-
fields : Optional[List[str]]
116+
fields : list[str] | None
115117
A list of fields to apply the action to. If 'None', the action is applied to the entire 'data'.
116118
action : Callable
117119
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(
142144

143145
def _apply_action_to_fields(
144146
self,
145-
data: Union[dict, str],
147+
data: dict | str,
146148
fields: list,
147149
action: Callable,
148150
provider_options: dict | None = None,
149151
**encryption_context: str,
150-
) -> Union[dict, str]:
152+
) -> dict | str:
151153
"""
152154
This method takes the input data, which can be either a dictionary or a JSON string,
153155
and erases, encrypts, or decrypts the specified fields.
154156
155157
Parameters
156158
----------
157-
data : Union[dict, str])
159+
data : dict | str)
158160
The input data to process. It can be either a dictionary or a JSON string.
159-
fields : List
161+
fields : list
160162
A list of fields to apply the action to. Each field can be specified as a string or
161163
a list of strings representing nested keys in the dictionary.
162164
action : Callable

aws_lambda_powertools/utilities/data_masking/provider/base.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ def encrypt(self, data) -> str:
2424
def decrypt(self, data) -> Any:
2525
# Implementation logic for data decryption
2626
27-
def erase(self, data) -> Union[str, Iterable]:
27+
def erase(self, data) -> str | Iterable:
2828
# Implementation logic for data masking
2929
pass
3030

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

+5-8
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import json
55
import logging
66
from binascii import Error
7-
from typing import Any, Callable, List
7+
from typing import Any, Callable
88

99
import botocore
1010
from aws_encryption_sdk import (
@@ -18,7 +18,6 @@
1818
GenerateKeyError,
1919
NotSupportedError,
2020
)
21-
from aws_encryption_sdk.structures import MessageHeader
2221

2322
from aws_lambda_powertools.shared.functions import (
2423
base64_decode,
@@ -76,7 +75,7 @@ def lambda_handler(event, context):
7675

7776
def __init__(
7877
self,
79-
keys: List[str],
78+
keys: list[str],
8079
key_provider=None,
8180
local_cache_capacity: int = CACHE_CAPACITY,
8281
max_cache_age_seconds: float = MAX_CACHE_AGE_SECONDS,
@@ -112,7 +111,7 @@ class KMSKeyProvider:
112111

113112
def __init__(
114113
self,
115-
keys: List[str],
114+
keys: list[str],
116115
json_serializer: Callable[..., str],
117116
json_deserializer: Callable[[str], Any],
118117
local_cache_capacity: int = CACHE_CAPACITY,
@@ -143,7 +142,7 @@ def encrypt(self, data: Any, provider_options: dict | None = None, **encryption_
143142
144143
Parameters
145144
-------
146-
data : Union[bytes, str]
145+
data : Any
147146
The data to be encrypted.
148147
provider_options : dict
149148
Additional options for the aws_encryption_sdk.EncryptionSDKClient
@@ -180,7 +179,7 @@ def decrypt(self, data: str, provider_options: dict | None = None, **encryption_
180179
181180
Parameters
182181
-------
183-
data : Union[bytes, str]
182+
data : str
184183
The encrypted data, as a base64-encoded string
185184
provider_options
186185
Additional options for the aws_encryption_sdk.EncryptionSDKClient
@@ -201,8 +200,6 @@ def decrypt(self, data: str, provider_options: dict | None = None, **encryption_
201200
)
202201

203202
try:
204-
decryptor_header: MessageHeader
205-
206203
ciphertext, decryptor_header = self.client.decrypt(
207204
source=ciphertext_decoded,
208205
key_provider=self.key_provider,

0 commit comments

Comments
 (0)