add return type hints for plot methods + fix up test assertions

Author: drbenvincent

causalpy/experiments/diff_in_diff.py

@@ -17,6 +17,7 @@
 
 import numpy as np
 import pandas as pd
+from matplotlib import pyplot as plt
 from patsy import build_design_matrices, dmatrices
 
 from causalpy.custom_exceptions import (
@@ -204,7 +205,7 @@ def input_validation(self):
                 coded. Consisting of 0's and 1's only."""
             )
 
-    def plot(self, round_to=None):
+    def plot(self, round_to=None) -> tuple[plt.Figure, plt.Axes]:
         """
         Plot the results
 

causalpy/experiments/inverse_propensity_weighting.py

@@ -15,6 +15,8 @@
 Inverse propensity weighting
 """
 
+from typing import List
+
 import arviz as az
 import matplotlib.pyplot as plt
 import numpy as np
@@ -238,7 +240,9 @@ def get_ate(self, i, idata, method="doubly_robust"):
         ate = trt - ntrt
         return [ate, trt, ntrt]
 
-    def plot_ate(self, idata=None, method=None, prop_draws=100, ate_draws=300):
+    def plot_ate(
+        self, idata=None, method=None, prop_draws=100, ate_draws=300
+    ) -> tuple[plt.Figure, List[plt.Axes]]:
         if idata is None:
             idata = self.model.idata
         if method is None:
@@ -364,7 +368,7 @@ def make_hists(idata, i, axs, method=method):
         axs[2].legend()
         axs[2].set_title("Average Treatment Effect", fontsize=20)
 
-        return fig
+        return fig, axs
 
     def weighted_percentile(self, data, weights, perc):
         """
@@ -380,7 +384,9 @@ def weighted_percentile(self, data, weights, perc):
         )  # 'like' a CDF function
         return np.interp(perc, cdf, data)
 
-    def plot_balance_ecdf(self, covariate, idata=None, weighting_scheme=None):
+    def plot_balance_ecdf(
+        self, covariate, idata=None, weighting_scheme=None
+    ) -> tuple[plt.Figure, List[plt.Axes]]:
         """
         Plotting function takes a single covariate and shows the
         differences in the ECDF between the treatment and control
@@ -451,4 +457,5 @@ def plot_balance_ecdf(self, covariate, idata=None, weighting_scheme=None):
         axs[0].set_xlabel("Quantiles")
         axs[1].legend()
         axs[0].legend()
-        return fig
+        # TODO: for some reason ax is type numpy.ndarray, so we need to convert this back to a list to conform to the expected return type.
+        return fig, list(axs)

causalpy/experiments/prepostfit.py

@@ -19,6 +19,7 @@
 
 import numpy as np
 import pandas as pd
+from matplotlib import pyplot as plt
 from patsy import build_design_matrices, dmatrices
 
 from causalpy.custom_exceptions import BadIndexException
@@ -107,7 +108,7 @@ def input_validation(self, data, treatment_time):
                 "If data.index is not DatetimeIndex, treatment_time must be pd.Timestamp."  # noqa: E501
             )
 
-    def plot(self):
+    def plot(self) -> tuple[plt.Figure, plt.Axes]:
         """
         Plot the results
 
@@ -204,7 +205,9 @@ class SyntheticControl(PrePostFit):
 
     expt_type = "SyntheticControl"
 
-    def plot(self, round_to=None, plot_predictors: bool = False):
+    def plot(
+        self, round_to=None, plot_predictors: bool = False
+    ) -> tuple[plt.Figure, plt.Axes]:
         """
         Plot the results
 

causalpy/experiments/prepostnegd.py

@@ -17,6 +17,7 @@
 
 import numpy as np
 import pandas as pd
+from matplotlib import pyplot as plt
 from patsy import build_design_matrices, dmatrices
 
 from causalpy.custom_exceptions import (
@@ -176,7 +177,7 @@ def _causal_impact_summary_stat(self, round_to) -> str:
         causal_impact = f"{round_num(self.causal_impact.mean(), round_to)}, "
         return f"Causal impact = {causal_impact + ci}"
 
-    def plot(self):
+    def plot(self) -> tuple[plt.Figure, plt.Axes]:
         """
         Plot the results
 

causalpy/experiments/regression_discontinuity.py

@@ -17,6 +17,7 @@
 
 import warnings  # noqa: I001
 
+from matplotlib import pyplot as plt
 import numpy as np
 import pandas as pd
 from patsy import build_design_matrices, dmatrices
@@ -192,7 +193,7 @@ def _is_treated(self, x):
         """
         return np.greater_equal(x, self.treatment_threshold)
 
-    def plot(self):
+    def plot(self) -> tuple[plt.Figure, plt.Axes]:
         """
         Plot the results
 

causalpy/experiments/regression_kink.py

@@ -18,6 +18,7 @@
 
 import warnings  # noqa: I001
 
+from matplotlib import pyplot as plt
 import numpy as np
 import pandas as pd
 from patsy import build_design_matrices, dmatrices
@@ -160,7 +161,7 @@ def _is_treated(self, x):
         """Returns ``True`` if `x` is greater than or equal to the treatment threshold."""  # noqa: E501
         return np.greater_equal(x, self.kink_point)
 
-    def plot(self, round_to=None):
+    def plot(self, round_to=None) -> tuple[plt.Figure, plt.Axes]:
         """
         Plot the results
 

causalpy/plotting.py

@@ -12,6 +12,8 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
+from typing import List
+
 import arviz as az
 import matplotlib.pyplot as plt
 import numpy as np
@@ -28,7 +30,9 @@ class BayesianPlotComponent:
     """Plotting component for Bayesian models."""
 
     @staticmethod
-    def plot_pre_post(results, round_to=None, counterfactual_label=None):
+    def plot_pre_post(
+        results, round_to=None, counterfactual_label=None
+    ) -> tuple[plt.Figure, plt.Axes]:
         """Generate plot for pre-post experiment types, such as Interrupted Time Series
         and Synthetic Control."""
         datapre = results.datapre
@@ -143,7 +147,9 @@ def plot_pre_post(results, round_to=None, counterfactual_label=None):
         return fig, ax
 
     @staticmethod
-    def plot_difference_in_differences(results, round_to=None):
+    def plot_difference_in_differences(
+        results, round_to=None
+    ) -> tuple[plt.Figure, plt.Axes]:
         """Generate plot for difference-in-differences"""
 
         def _plot_causal_impact_arrow(results, ax):
@@ -282,7 +288,7 @@ def _plot_causal_impact_arrow(results, ax):
     #     pass
 
     @staticmethod
-    def plot_pre_post_negd(results, round_to=None):
+    def plot_pre_post_negd(results, round_to=None) -> tuple[plt.Figure, List[plt.Axes]]:
         """Generate plot for ANOVA-like experiments with non-equivalent group designs."""
         fig, ax = plt.subplots(
             2, 1, figsize=(7, 9), gridspec_kw={"height_ratios": [3, 1]}
@@ -334,7 +340,9 @@ def plot_pre_post_negd(results, round_to=None):
         return fig, ax
 
     @staticmethod
-    def plot_regression_discontinuity(results, round_to=None):
+    def plot_regression_discontinuity(
+        results, round_to=None
+    ) -> tuple[plt.Figure, plt.Axes]:
         """Generate plot for regression discontinuity designs."""
         fig, ax = plt.subplots()
         # Plot raw data
@@ -386,7 +394,7 @@ def plot_regression_discontinuity(results, round_to=None):
         return (fig, ax)
 
     @staticmethod
-    def plot_regression_kink(results, round_to=None):
+    def plot_regression_kink(results, round_to=None) -> tuple[plt.Figure, plt.Axes]:
         """Generate plot for regression kink designs."""
         fig, ax = plt.subplots()
         # Plot raw data
@@ -440,7 +448,7 @@ class OLSPlotComponent:
     """Plotting component for OLS models."""
 
     @staticmethod
-    def plot_pre_post(results, round_to=None):
+    def plot_pre_post(results, round_to=None) -> tuple[plt.Figure, List[plt.Axes]]:
         """Generate plot for pre-post experiment types, such as Interrupted Time Series
         and Synthetic Control."""
         counterfactual_label = "Counterfactual"
@@ -509,7 +517,9 @@ def plot_pre_post(results, round_to=None):
         return (fig, ax)
 
     @staticmethod
-    def plot_difference_in_differences(results, round_to=None):
+    def plot_difference_in_differences(
+        results, round_to=None
+    ) -> tuple[plt.Figure, plt.Axes]:
         """Generate plot for difference-in-differences"""
         fig, ax = plt.subplots()
 
@@ -585,7 +595,9 @@ def plot_difference_in_differences(results, round_to=None):
         return (fig, ax)
 
     @staticmethod
-    def plot_regression_discontinuity(results, round_to=None) -> tuple:
+    def plot_regression_discontinuity(
+        results, round_to=None
+    ) -> tuple[plt.Figure, plt.Axes]:
         """Generate plot for regression discontinuity designs."""
         fig, ax = plt.subplots()
         # Plot raw data

causalpy/tests/test_misc.py

@@ -16,8 +16,8 @@
 """
 
 import arviz as az
-import matplotlib as mpl
 import pandas as pd
+from matplotlib import pyplot as plt
 
 import causalpy as cp
 
@@ -82,7 +82,11 @@ def test_inverse_prop():
     assert isinstance(ate_list, list)
     ate_list = result.get_ate(0, result.idata, method="overlap")
     assert isinstance(ate_list, list)
-    fig = result.plot_ate(prop_draws=1, ate_draws=10)
-    assert isinstance(fig, mpl.figure.Figure)
-    fig = result.plot_balance_ecdf("age")
-    assert isinstance(fig, mpl.figure.Figure)
+    fig, axs = result.plot_ate(prop_draws=1, ate_draws=10)
+    assert isinstance(fig, plt.Figure)
+    assert isinstance(axs, list)
+    assert all(isinstance(ax, plt.Axes) for ax in axs)
+    fig, axs = result.plot_balance_ecdf("age")
+    assert isinstance(fig, plt.Figure)
+    assert isinstance(axs, list)
+    assert all(isinstance(ax, plt.Axes) for ax in axs)