Add table method to ModelBuilder to display rich table (#1786)

Author: williambdean

pymc_marketing/model_builder.py

@@ -17,6 +17,7 @@
 import json
 import warnings
 from abc import ABC, abstractmethod
+from functools import wraps
 from inspect import signature
 from pathlib import Path
 from typing import Any
@@ -27,6 +28,8 @@
 import pymc as pm
 import xarray as xr
 from pymc.util import RandomState
+from pymc_extras.printing import model_table
+from rich.table import Table
 
 from pymc_marketing.hsgp_kwargs import HSGPKwargs
 from pymc_marketing.utils import from_netcdf
@@ -103,6 +106,20 @@ def accessor(self) -> xr.Dataset:
     return property(accessor)
 
 
+def requires_model(func):
+    """Ensure that the model is built before calling a method."""
+
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        if not hasattr(self, "model"):
+            raise RuntimeError(
+                "The model hasn't been built yet. Please call `build_model` first."
+            )
+        return func(self, *args, **kwargs)
+
+    return wrapper
+
+
 def create_sample_kwargs(
     sampler_config: dict[str, Any],
     progressbar: bool | None,
@@ -1063,6 +1080,23 @@ def graphviz(self, **kwargs):
         """
         return pm.model_to_graphviz(self.model, **kwargs)
 
+    @requires_model
+    def table(self, **model_table_kwargs) -> Table:
+        """Get the summary table of the model.
+
+        Parameters
+        ----------
+        **model_table_kwargs
+            Keyword arguments for the `model_table` function
+
+        Returns
+        -------
+        rich.table.Table
+            A rich table containing the summary of the model.
+
+        """
+        return model_table(self.model, **model_table_kwargs)
+
     prior = create_idata_accessor(
         "prior",
         "The model hasn't been sampled yet, call .sample_prior_predictive() first",

tests/test_model_builder.py

@@ -24,6 +24,7 @@
 import pymc as pm
 import pytest
 import xarray as xr
+from rich.table import Table
 
 from pymc_marketing.model_builder import (
     DifferentModelError,
@@ -794,3 +795,13 @@ def test_load_from_idata_errors(request, fixture_name, match) -> None:
     idata = request.getfixturevalue(fixture_name)
     with pytest.raises(DifferentModelError, match=match):
         ModelBuilderTest.load_from_idata(idata)
+
+
+def test_table() -> None:
+    model = ModelBuilderTest()
+    match = "The model hasn't been built yet"
+    with pytest.raises(RuntimeError, match=match):
+        model.table()
+
+    model.build_model(pd.DataFrame({"input": [1, 2, 3]}), pd.Series([1, 2, 3]))
+    assert isinstance(model.table(), Table)