tests

Author: juanitorduz

pymc_marketing/mmm/mmm.py

@@ -857,7 +857,7 @@ def build_model(  # type: ignore[override]
                 raise TypeError("X data must be a DataFrame")
 
             y_data = self.preprocessed_data["y"]
-            if not isinstance(y_data, (pd.Series, np.ndarray)):
+            if not isinstance(y_data, pd.Series | np.ndarray):
                 raise TypeError("y data must be a Series or ndarray")
 
             channel_data_scaled, target_scaled, _, target_scale_ = (
@@ -1154,7 +1154,7 @@ def _prepare_target_data(
         if y is None:
             # When y is None, create zeros array matching the type of preprocessed y
             y_preprocessed = self.preprocessed_data["y"]
-            if isinstance(y_preprocessed, (pd.Series, np.ndarray)):
+            if isinstance(y_preprocessed, pd.Series | np.ndarray):
                 y_data = np.zeros(n_rows, dtype=np.asarray(y_preprocessed).dtype)
             else:
                 # Default to float64 if type is unknown

tests/mmm/test_mmm.py

@@ -2050,7 +2050,7 @@ def test_compute_scales(self, toy_X: pd.DataFrame, toy_y: pd.Series):
         assert hasattr(mmm, "channel_scale")
         assert hasattr(mmm, "target_scale")
         assert len(mmm.channel_scale) == len(mmm.channel_columns)
-        assert isinstance(mmm.target_scale, (float, np.floating))
+        assert isinstance(mmm.target_scale, float | np.floating)
 
     def test_build_intercept_static(self, toy_X: pd.DataFrame, toy_y: pd.Series):
         """Test _build_intercept with static intercept."""