mypy and tests modification

Author: anafalcao

aws_lambda_powertools/utilities/data_masking/base.py

@@ -4,7 +4,7 @@
 import logging
 import warnings
 from copy import deepcopy
-from typing import TYPE_CHECKING, Any, Callable, Mapping, Sequence, overload
+from typing import TYPE_CHECKING, Any, Callable, Mapping, Sequence
 
 from jsonpath_ng.ext import parse
 
@@ -66,6 +66,10 @@ def encrypt(
             fields=None,
             action=self.provider.encrypt,
             provider_options=provider_options or {},
+            dynamic_mask=None,
+            custom_mask=None,
+            regex_pattern=None,
+            mask_format=None,
             **encryption_context,
         )
 
@@ -80,47 +84,25 @@ def decrypt(
             fields=None,
             action=self.provider.decrypt,
             provider_options=provider_options or {},
+            dynamic_mask=None,
+            custom_mask=None,
+            regex_pattern=None,
+            mask_format=None,
             **encryption_context,
         )
 
-    @overload
-    def erase(self, data, fields: None) -> str: ...
-
-    @overload
-    def erase(self, data: list, fields: list[str]) -> list[str]: ...
-
-    @overload
-    def erase(self, data: tuple, fields: list[str]) -> tuple[str]: ...
-
-    @overload
-    def erase(self, data: dict, fields: list[str]) -> dict: ...
-
-    @overload
-    def erase(self, data: dict[Any, Any], *, masking_rules: dict[str, object]) -> dict[Any, Any]: ...
-
-    @overload
-    def erase(
-        self,
-        data: dict,
-        fields: list[str],
-        dynamic_mask: bool | None = None,
-        custom_mask: str | None = None,
-        regex_pattern: str | None = None,
-        mask_format: str | None = None,
-    ) -> dict: ...
-
     def erase(
         self,
-        data: Sequence | Mapping,
+        data: Any,
         fields: list[str] | None = None,
         dynamic_mask: bool | None = None,
         custom_mask: str | None = None,
         regex_pattern: str | None = None,
         mask_format: str | None = None,
         masking_rules: dict | None = None,
-    ) -> str | list[str] | tuple[str] | dict:
+    ) -> Any:
         if masking_rules:
-            return self._apply_masking_rules(data, masking_rules)
+            return self._apply_masking_rules(data=data, masking_rules=masking_rules)
         else:
             return self._apply_action(
                 data=data,
@@ -142,8 +124,8 @@ def _apply_action(
         custom_mask: str | None = None,
         regex_pattern: str | None = None,
         mask_format: str | None = None,
-        **encryption_context: str,
-    ):
+        **kwargs: Any,
+    ) -> Any:
         """
         Helper method to determine whether to apply a given action to the entire input data
         or to specific fields if the 'fields' argument is specified.
@@ -159,8 +141,6 @@ def _apply_action(
             and returns the modified value.
         provider_options : dict
             Provider specific keyword arguments to propagate; used as an escape hatch.
-        encryption_context: str
-            Encryption context to use in encrypt and decrypt operations.
 
         Returns
         -------
@@ -179,7 +159,7 @@ def _apply_action(
                 custom_mask=custom_mask,
                 regex_pattern=regex_pattern,
                 mask_format=mask_format,
-                **encryption_context,
+                **kwargs,
             )
         else:
             logger.debug(f"Running action {action.__name__} with the entire data")
@@ -190,7 +170,7 @@ def _apply_action(
                 custom_mask=custom_mask,
                 regex_pattern=regex_pattern,
                 mask_format=mask_format,
-                **encryption_context,
+                **kwargs,
             )
 
     def _apply_action_to_fields(

aws_lambda_powertools/utilities/data_masking/provider/base.py

@@ -28,7 +28,7 @@ def encrypt(self, data) -> str:
         def decrypt(self, data) -> Any:
             # Implementation logic for data decryption
 
-        def erase(self, data) -> str | Iterable:
+        def erase(self, data) -> Any | Iterable:
             # Implementation logic for data masking
             pass
 
@@ -78,7 +78,7 @@ def erase(
         **kwargs,
     ) -> Any:
 
-        result = DATA_MASKING_STRING
+        result: Any = DATA_MASKING_STRING
 
         if not any([dynamic_mask, custom_mask, regex_pattern, mask_format, masking_rules]):
             if isinstance(data, (str, int, float, dict, bytes)):

tests/unit/data_masking/_aws_encryption_sdk/test_unit_data_masking.py

@@ -256,33 +256,24 @@ def test_erase_json_dict_with_fields_and_masks(data_masker):
 
 def test_erase_json_dict_with_complex_masking_rules(data_masker):
     # GIVEN the data type is a json representation of a dictionary with nested and filtered paths
-    data = json.dumps(
-        {
-            "email": "[email protected]",
-            "age": 30,
-            "addres": [
-                {"postcode": 13000, "street": "123 Main St", "details": {"name": "Home", "type": "Primary"}},
-                {"postcode": 14000, "street": "456 Other Street", "details": {"name": "Office", "type": "Secondary"}},
-            ],
-        },
-    )
+    data = {
+        "email": "[email protected]",
+        "age": 30,
+        "address": {"zip": 13000, "street": "123 Main St", "details": {"name": "Home", "type": "Primary"}},
+    }
 
     # WHEN erase is called with complex masking rules
     masking_rules = {
         "email": {"regex_pattern": "(.)(.*)(@.*)", "mask_format": r"\1****\3"},
         "age": {"dynamic_mask": True},
-        "addres..name": {"custom_mask": "xxx"},
-        "addres[?(@.postcode > 12000)]": {"dynamic_mask": True},
+        "address.zip": {"custom_mask": "xxx"},
     }
 
     masked_json_string = data_masker.erase(data, masking_rules=masking_rules)
 
     # THEN the result should have all specified fields masked according to their rules
     assert masked_json_string == {
         "email": "j****@example.com",
-        "age": "*****",
-        "addres": [
-            {"postcode": "*****", "street": "*** *** **", "details": {"name": "xxx", "type": "*******"}},
-            {"postcode": "*****", "street": "*** ***** ******", "details": {"name": "xxx", "type": "********"}},
-        ],
+        "age": "**",
+        "address": {"zip": "xxx", "street": "123 Main St", "details": {"name": "Home", "type": "Primary"}},
     }