swap assert for a custom BadIndexException defined in a custom_exceptions module

Author: drbenvincent

causalpy/custom_exceptions.py

@@ -0,0 +1,6 @@
+class BadIndexException(Exception):
+    """Custom exception used when we have a mismatch in types between the dataframe
+    index and an event, typically a treatment or intervention."""
+
+    def __init__(self, message):
+        self.message = message

causalpy/pymc_experiments.py

@@ -8,6 +8,7 @@
 import xarray as xr
 from patsy import build_design_matrices, dmatrices
 
+from causalpy.custom_exceptions import BadIndexException
 from causalpy.plot_utils import plot_xY
 
 LEGEND_FONT_SIZE = 12
@@ -117,14 +118,18 @@ def __init__(
 
     def _input_validation(self, data, treatment_time):
         """Validate the input data for correctness"""
-        if isinstance(data.index, pd.DatetimeIndex):
-            assert isinstance(
-                treatment_time, pd.Timestamp
-            ), "If data.index is DatetimeIndex, treatment_time must be pd.Timestamp."
-        else:
-            assert (
-                isinstance(treatment_time, pd.Timestamp) is False
-            ), "If treatment_time is pd.Timestamp, this only makese sense if data.index is DatetimeIndex."  # noqa: E501
+        if isinstance(data.index, pd.DatetimeIndex) and not isinstance(
+            treatment_time, pd.Timestamp
+        ):
+            raise BadIndexException(
+                "If data.index is DatetimeIndex, treatment_time must be pd.Timestamp."
+            )
+        if not isinstance(data.index, pd.DatetimeIndex) and isinstance(
+            treatment_time, pd.Timestamp
+        ):
+            raise BadIndexException(
+                "If data.index is not DatetimeIndex, treatment_time must be pd.Timestamp."  # noqa: E501
+            )
 
     def plot(self):
 

causalpy/tests/test_integration_pymc_examples.py

@@ -2,6 +2,7 @@
 import pytest
 
 import causalpy as cp
+from causalpy.custom_exceptions import BadIndexException
 
 sample_kwargs = {"tune": 20, "draws": 20, "chains": 2, "cores": 2}
 
@@ -193,9 +194,9 @@ def test_sc():
 
 @pytest.mark.integration
 def test_sc_input_error():
-    """Test that an error is raised if the data index is not datetime and the
-    treatment time is pd.Timestamp."""
-    with pytest.raises(AssertionError):
+    """Confirm that a BadIndexException is raised treatment_time is pd.Timestamp
+    and df.index is not pd.DatetimeIndex."""
+    with pytest.raises(BadIndexException):
         df = cp.load_data("sc")
         treatment_time = pd.to_datetime("2016 June 24")
         _ = cp.pymc_experiments.SyntheticControl(
@@ -234,9 +235,9 @@ def test_sc_brexit():
 
 @pytest.mark.integration
 def test_sc_brexit_input_error():
-    """Test that an error is raised if the data index is datetime and the treatment time
-    is not pd.Timestamp."""
-    with pytest.raises(AssertionError):
+    """Confirm a BadIndexException is raised if the data index is datetime and the
+    treatment time is not pd.Timestamp."""
+    with pytest.raises(BadIndexException):
         df = cp.load_data("brexit")
         df["Time"] = pd.to_datetime(df["Time"])
         df.set_index("Time", inplace=True)