Merge branch 'main' into issue-2661-download-demo-update

Author: fealho

sdv/cag/_utils.py

@@ -5,7 +5,7 @@
 import pandas as pd
 
 from sdv.cag._errors import ConstraintNotMetError
-from sdv.errors import SynthesizerInputError, TableNameError
+from sdv.errors import RefitWarning, SynthesizerInputError, TableNameError
 from sdv.metadata import Metadata
 
 
@@ -185,7 +185,8 @@ def _validate_constraints(constraints, synthesizer_fitted):
 
     if synthesizer_fitted:
         warnings.warn(
-            "For these constraints to take effect, please refit the synthesizer using 'fit'."
+            "For these constraints to take effect, please refit the synthesizer using 'fit'.",
+            RefitWarning,
         )
 
     return _filter_old_style_constraints(constraints)

sdv/errors.py

@@ -91,3 +91,11 @@ class DemoResourceNotFoundError(Exception):
     This error is intended for missing demo assets such as the dataset archive,
     metadata, license, README, or other auxiliary files in the demo bucket.
     """
+
+
+class RefitWarning(UserWarning):
+    """Warning to be raised if the synthesizer needs to be refit.
+
+    Warning to be raised if a change to a synthesizer requires the synthesizer
+    to be refit for the change to be applied.
+    """

sdv/multi_table/base.py

@@ -29,6 +29,7 @@
 from sdv.cag.programmable_constraint import ProgrammableConstraint, ProgrammableConstraintHarness
 from sdv.errors import (
     InvalidDataError,
+    RefitWarning,
     SamplingError,
     SynthesizerInputError,
 )
@@ -551,10 +552,11 @@ def preprocess(self, data):
         self.validate(data)
         data = self._validate_transform_constraints(data)
         if self._fitted:
-            warnings.warn(
+            msg = (
                 'This model has already been fitted. To use the new preprocessed data, '
                 "please refit the model using 'fit' or 'fit_processed_data'."
             )
+            warnings.warn(msg, RefitWarning)
 
         processed_data = {}
         pbar_args = self._get_pbar_args(desc='Preprocess Tables')

sdv/single_table/base.py

@@ -41,6 +41,7 @@
 from sdv.errors import (
     ConstraintsNotMetError,
     InvalidDataError,
+    RefitWarning,
     SamplingError,
     SynthesizerInputError,
 )
@@ -306,7 +307,7 @@ def update_transformers(self, column_name_to_transformer):
         self._data_processor.update_transformers(column_name_to_transformer)
         if self._fitted:
             msg = 'For this change to take effect, please refit the synthesizer using `fit`.'
-            warnings.warn(msg, UserWarning)
+            warnings.warn(msg, RefitWarning)
 
     def get_parameters(self):
         """Return the parameters used to instantiate the synthesizer."""
@@ -587,10 +588,12 @@ def _preprocess_helper(self, data):
         """
         self.validate(data)
         if self._fitted:
-            warnings.warn(
+            msg = (
                 'This model has already been fitted. To use the new preprocessed data, '
                 "please refit the model using 'fit' or 'fit_processed_data'."
             )
+            warnings.warn(msg, RefitWarning)
+
         data = self._validate_transform_constraints(data)
 
         return data

tests/unit/multi_table/test_base.py

@@ -17,6 +17,7 @@
 from sdv.errors import (
     InvalidDataError,
     NotFittedError,
+    RefitWarning,
     SamplingError,
     SynthesizerInputError,
     VersionError,
@@ -1013,7 +1014,8 @@ def test_preprocess_warning(self, mock_warnings):
         assert args[0].equals(data['upravna_enota'])
         mock_warnings.warn.assert_called_once_with(
             'This model has already been fitted. To use the new preprocessed data, '
-            "please refit the model using 'fit' or 'fit_processed_data'."
+            "please refit the model using 'fit' or 'fit_processed_data'.",
+            RefitWarning,
         )
 
     @patch('sdv.metadata.single_table.SingleTableMetadata._validate_metadata_matches_data')

tests/unit/single_table/test_base.py

@@ -23,6 +23,7 @@
 from sdv.errors import (
     ConstraintsNotMetError,
     InvalidDataError,
+    RefitWarning,
     SamplingError,
     SynthesizerInputError,
     VersionError,
@@ -665,7 +666,7 @@ def test__preprocess_helper(self, mock_warnings):
         result = BaseSynthesizer._preprocess_helper(instance, data)
 
         # Assert
-        mock_warnings.warn.assert_called_once_with(expected_warning)
+        mock_warnings.warn.assert_called_once_with(expected_warning, RefitWarning)
         instance.validate.assert_called_once_with(data)
         instance._validate_transform_constraints.assert_called_once_with(data)
         pd.testing.assert_frame_equal(result, data)