move out _convert_original_cols from _preprocess_helper

Author: R-Palazzo

sdv/single_table/base.py

@@ -547,15 +547,14 @@ def _preprocess_helper(self, data):
         - Warn the user if the model has already been fitted
         - Store the original columns and convert them to string if needed
         """
-        is_converted = self._store_and_convert_original_cols(data)
         self.validate(data)
         if self._fitted:
             warnings.warn(
                 'This model has already been fitted. To use the new preprocessed data, '
                 "please refit the model using 'fit' or 'fit_processed_data'."
             )
 
-        return data, is_converted
+        return data
 
     def preprocess(self, data):
         """Transform the raw data to numerical space.
@@ -568,7 +567,8 @@ def preprocess(self, data):
             pandas.DataFrame:
                 The preprocessed data.
         """
-        data, is_converted = self._preprocess_helper(data)
+        is_converted = self._store_and_convert_original_cols(data)
+        data = self._preprocess_helper(data)
         preprocess_data = self._preprocess(data)
         if is_converted:
             data.columns = self._original_columns
@@ -856,10 +856,10 @@ def validate(self, data):
         self.metadata = metadata
 
     def _preprocess_helper(self, data):
-        data, is_converted = super()._preprocess_helper(data)
+        data = super()._preprocess_helper(data)
         data = self._transform_helper(data)
 
-        return data, is_converted
+        return data
 
     def _set_random_state(self, random_state):
         """Set the random state of the model's random number generator.

tests/unit/single_table/test_base.py

@@ -505,7 +505,6 @@ def test__preprocess_helper_basesynthesizer(self, mock_warnings):
         instance = Mock()
         instance._fitted = True
         data = pd.DataFrame({'name': ['John', 'Doe', 'John Doe']})
-        instance._store_and_convert_original_cols.return_value = False
 
         # Run
         result = BaseSynthesizer._preprocess_helper(instance, data)
@@ -517,8 +516,7 @@ def test__preprocess_helper_basesynthesizer(self, mock_warnings):
         )
         mock_warnings.warn.assert_called_once_with(expected_warning)
         instance.validate.assert_called_once_with(data)
-        pd.testing.assert_frame_equal(result[0], data)
-        assert result[1] is False
+        pd.testing.assert_frame_equal(result, data)
 
     @patch.object(BaseSynthesizer, '_preprocess_helper')
     def test__preprocess_helper(self, mock_preprocess_helper):
@@ -536,14 +534,13 @@ def test__preprocess_helper(self, mock_preprocess_helper):
         instance = BaseSingleTableSynthesizer(metadata)
         data = pd.DataFrame({'name': ['John', 'Doe', 'John Doe']})
         instance._transform_helper = Mock(return_value=data)
-        mock_preprocess_helper.return_value = (data, False)
+        mock_preprocess_helper.return_value = data
 
         # Run
         result = instance._preprocess_helper(data)
 
         # Assert
-        pd.testing.assert_frame_equal(result[0], data)
-        assert result[1] is False
+        pd.testing.assert_frame_equal(result, data)
         instance._transform_helper.assert_called_once_with(data)
 
     def test__preprocess(self):
@@ -569,22 +566,20 @@ def test__preprocess(self):
         )
 
     def test_preprocess(self):
-        """Test the preprocess method.
-
-        The preprocess method raises a warning if it was already fitted and then calls
-        ``_preprocess``.
-        """
+        """Test the preprocess method."""
         # Setup
         instance = Mock()
         instance._fitted = True
         data = pd.DataFrame({'name': ['John', 'Doe', 'John Doe']})
-        instance._preprocess_helper.return_value = (data, False)
+        instance._preprocess_helper.return_value = data
+        instance._store_and_convert_original_cols = Mock(return_value=False)
         instance._preprocess.return_value = data
 
         # Run
         result = BaseSingleTableSynthesizer.preprocess(instance, data)
 
         # Assert
+        instance._store_and_convert_original_cols.assert_called_once_with(data)
         instance._preprocess_helper.assert_called_once_with(data)
         instance._preprocess.assert_called_once_with(data)
         pd.testing.assert_frame_equal(result, data)