Improve logic on the locale detection

sdv-dev · Sep 6, 2023 · 7486dbc · 7486dbc
1 parent b09f5a9
commit 7486dbc
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/rdt/transformers/pii/anonymizer.py b/rdt/transformers/pii/anonymizer.py
@@ -80,8 +80,12 @@ def _check_locales(self):
         locales = self.locales if isinstance(self.locales, list) else [self.locales]
         missed_locales = []
         for locale in locales:
-            spec = importlib.util.find_spec(f'faker.providers.{self.provider_name}.{locale}')
-            if spec is None and not self.provider_name.endswith(f'.{locale}'):
+            provider_name = self.provider_name
+            if self.provider_name.endswith(f'.{locale}'):
+                provider_name = self.provider_name.replace(f'.{locale}', '')
+
+            spec = importlib.util.find_spec(f'faker.providers.{provider_name}.{locale}')
+            if spec is None:
                 missed_locales.append(locale)
 
         if missed_locales:

diff --git a/tests/unit/transformers/pii/test_anonymizer.py b/tests/unit/transformers/pii/test_anonymizer.py
@@ -200,7 +200,7 @@ def test__check_locales_provider_ending_with_locale(self, mock_warnings, mock_im
         instance.provider_name = 'address.en_US'
         instance.function_name = 'postcode'
         instance.locales = ['en_US']
-        mock_importlib.util.find_spec.side_effect = [None]
+        mock_importlib.util.find_spec.side_effect = ['en_US']
 
         # Run
         AnonymizedFaker._check_locales(instance)