Address feedback

sdv-dev · Oct 17, 2023 · fa08380 · fa08380
1 parent 1ac0a8f
commit fa08380
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 10 deletions.
diff --git a/rdt/transformers/numerical.py b/rdt/transformers/numerical.py
@@ -272,6 +272,11 @@ class GaussianNormalizer(FloatFormatter):
     """
 
     _univariate = None
+    _DEPRECATED_DISTRIBUTIONS_MAPPING = {
+        'gaussian': 'norm',
+        'student_t': 't',
+        'truncated_gaussian': 'truncnorm'
+    }
 
     @staticmethod
     def _get_distributions():
@@ -292,10 +297,6 @@ def _get_distributions():
             'gaussian_kde': univariate.GaussianKDE,
             'truncnorm': univariate.TruncatedGaussian,
             'uniform': univariate.UniformUnivariate,
-            # the following are deprecated
-            'gaussian': univariate.GaussianUnivariate,
-            'student_t': univariate.StudentTUnivariate,
-            'truncated_gaussian': univariate.TruncatedGaussian,
         }
 
     def __init__(self, model_missing_values=None, learn_rounding_scheme=False,
@@ -316,13 +317,13 @@ def __init__(self, model_missing_values=None, learn_rounding_scheme=False,
         self._distributions = self._get_distributions()
         if isinstance(distribution, str):
             if distribution in {'gaussian', 'student_t', 'truncated_gaussian'}:
-                deprecated_distributions_mapping = {
-                    'gaussian': 'norm', 'student_t': 't', 'truncated_gaussian': 'truncnorm'}
                 warnings.warn(
                     f"Future versions of RDT will not support '{distribution}' as an option. "
-                    f"Please use '{deprecated_distributions_mapping[distribution]}' instead.",
+                    f"Please use '{self._DEPRECATED_DISTRIBUTIONS_MAPPING[distribution]}' "
+                    'instead.',
                     FutureWarning
                 )
+                distribution = self._DEPRECATED_DISTRIBUTIONS_MAPPING[distribution]
 
             distribution = self._distributions[distribution]
 

diff --git a/tests/unit/transformers/test_numerical.py b/tests/unit/transformers/test_numerical.py
@@ -849,12 +849,9 @@ def test__get_distributions(self):
 
         # Assert
         expected = {
-            'gaussian': univariate.GaussianUnivariate,
             'gamma': univariate.GammaUnivariate,
             'beta': univariate.BetaUnivariate,
-            'student_t': univariate.StudentTUnivariate,
             'gaussian_kde': univariate.GaussianKDE,
-            'truncated_gaussian': univariate.TruncatedGaussian,
             'uniform': univariate.UniformUnivariate,
             'truncnorm': univariate.TruncatedGaussian,
             'norm': univariate.GaussianUnivariate,