Skip to content

Commit bf15d57

Browse files
refactor(data_masking): use standard collections for types (#6493)
Using generics types
1 parent acfa2a5 commit bf15d57

File tree

9 files changed

+36
-8
lines changed

9 files changed

+36
-8
lines changed

aws_lambda_powertools/utilities/data_masking/base.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
import logging
1212
import warnings
1313
from copy import deepcopy
14-
from typing import TYPE_CHECKING, Any, Callable, Mapping, Sequence
14+
from typing import TYPE_CHECKING, Any
1515

1616
from jsonpath_ng.ext import parse
1717

@@ -23,6 +23,7 @@
2323
from aws_lambda_powertools.warnings import PowertoolsUserWarning
2424

2525
if TYPE_CHECKING:
26+
from collections.abc import Callable, Mapping, Sequence
2627
from numbers import Number
2728

2829
logger = logging.getLogger(__name__)

aws_lambda_powertools/utilities/data_masking/provider/base.py

+7-2
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,18 @@
33
import functools
44
import json
55
import re
6-
from typing import Any, Callable
6+
from typing import TYPE_CHECKING, Any
77

88
from aws_lambda_powertools.utilities.data_masking.constants import DATA_MASKING_STRING
99

10+
if TYPE_CHECKING:
11+
from collections.abc import Callable
12+
1013
PRESERVE_CHARS = set("-_. ")
1114
_regex_cache = {}
1215

16+
JSON_DUMPS_CALL = functools.partial(json.dumps, ensure_ascii=False)
17+
1318

1419
class BaseProvider:
1520
"""
@@ -49,7 +54,7 @@ def lambda_handler(event, context):
4954

5055
def __init__(
5156
self,
52-
json_serializer: Callable[..., str] = functools.partial(json.dumps, ensure_ascii=False),
57+
json_serializer: Callable[..., str] = JSON_DUMPS_CALL,
5358
json_deserializer: Callable[[str], Any] = json.loads,
5459
) -> None:
5560
self.json_serializer = json_serializer

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

+7-2
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
7+
from typing import TYPE_CHECKING, Any
88

99
import botocore
1010
from aws_encryption_sdk import (
@@ -41,8 +41,13 @@
4141
)
4242
from aws_lambda_powertools.utilities.data_masking.provider import BaseProvider
4343

44+
if TYPE_CHECKING:
45+
from collections.abc import Callable
46+
4447
logger = logging.getLogger(__name__)
4548

49+
JSON_DUMPS_CALL = functools.partial(json.dumps, ensure_ascii=False)
50+
4651

4752
class AWSEncryptionSDKProvider(BaseProvider):
4853
"""
@@ -81,7 +86,7 @@ def __init__(
8186
max_cache_age_seconds: float = MAX_CACHE_AGE_SECONDS,
8287
max_messages_encrypted: int = MAX_MESSAGES_ENCRYPTED,
8388
max_bytes_encrypted: int = MAX_BYTES_ENCRYPTED,
84-
json_serializer: Callable[..., str] = functools.partial(json.dumps, ensure_ascii=False),
89+
json_serializer: Callable[..., str] = JSON_DUMPS_CALL,
8590
json_deserializer: Callable[[str], Any] = json.loads,
8691
):
8792
super().__init__(json_serializer=json_serializer, json_deserializer=json_deserializer)

tests/functional/data_masking/_aws_encryption_sdk/test_aws_encryption_sdk.py

+10-3
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
1+
from __future__ import annotations
2+
13
import base64
24
import functools
35
import json
4-
from typing import Any, Callable, Union
6+
from typing import TYPE_CHECKING, Any
57

68
import pytest
79
from aws_encryption_sdk.identifiers import Algorithm
@@ -13,16 +15,21 @@
1315
AWSEncryptionSDKProvider,
1416
)
1517

18+
if TYPE_CHECKING:
19+
from collections.abc import Callable
20+
21+
JSON_DUMPS_CALL = functools.partial(json.dumps, ensure_ascii=False)
22+
1623

1724
class FakeEncryptionKeyProvider(BaseProvider):
1825
def __init__(
1926
self,
20-
json_serializer: Callable = functools.partial(json.dumps, ensure_ascii=False),
27+
json_serializer: Callable = JSON_DUMPS_CALL,
2128
json_deserializer: Callable = json.loads,
2229
) -> None:
2330
super().__init__(json_serializer, json_deserializer)
2431

25-
def encrypt(self, data: Union[bytes, str], **kwargs) -> str:
32+
def encrypt(self, data: bytes | str, **kwargs) -> str:
2633
encoded_data: str = self.json_serializer(data)
2734
ciphertext = base64.b64encode(encoded_data.encode("utf-8")).decode()
2835
return ciphertext

tests/functional/data_masking/_pydantic/test_data_masking_with_pydantic.py

+2
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from __future__ import annotations
2+
13
import dataclasses
24

35
import pytest

tests/functional/data_masking/conftest.py

+2
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from __future__ import annotations
2+
13
from pytest_socket import disable_socket
24

35

tests/functional/data_masking/required_dependencies/test_erase_data_masking.py

+2
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from __future__ import annotations
2+
13
import json
24

35
import pytest

tests/unit/data_masking/_aws_encryption_sdk/test_kms_provider.py

+2
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from __future__ import annotations
2+
13
import pytest
24

35
from aws_lambda_powertools.utilities.data_masking.exceptions import (

tests/unit/data_masking/required_dependencies/test_base_functions.py

+2
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from __future__ import annotations
2+
13
import pytest
24

35
from aws_lambda_powertools.utilities.data_masking.base import DataMasking

0 commit comments

Comments
 (0)