Do not monkey-patch Ipython pretty representation on model variables

Author: ricardoV94

pymc/distributions/distribution.py

@@ -12,10 +12,8 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 import contextvars
-import functools
 import re
 import sys
-import types
 import warnings
 
 from abc import ABCMeta
@@ -53,7 +51,6 @@
 from pymc.logprob.abstract import MeasurableOp, _icdf, _logcdf, _logprob
 from pymc.logprob.basic import logp
 from pymc.logprob.rewriting import logprob_rewrites_db
-from pymc.printing import str_for_dist
 from pymc.pytensorf import (
     collect_default_updates_inner_fgraph,
     constant_fold,
@@ -514,12 +511,6 @@ def __new__(
             default_transform=default_transform,
             initval=initval,
         )
-
-        # add in pretty-printing support
-        rv_out.str_repr = types.MethodType(str_for_dist, rv_out)
-        rv_out._repr_latex_ = types.MethodType(
-            functools.partial(str_for_dist, formatting="latex"), rv_out
-        )
         return rv_out
 
     @classmethod

pymc/model/core.py

@@ -16,7 +16,6 @@
 import functools
 import sys
 import threading
-import types
 import warnings
 
 from collections.abc import Iterable, Sequence
@@ -491,13 +490,6 @@ def __init__(
             self._dim_lengths = {}
         self.add_coords(coords)
 
-        from pymc.printing import str_for_model
-
-        self.str_repr = types.MethodType(str_for_model, self)
-        self._repr_latex_ = types.MethodType(
-            functools.partial(str_for_model, formatting="latex"), self
-        )
-
     @classmethod
     def get_context(
         cls, error_if_none: bool = True, allow_block_model_access: bool = False
@@ -2014,6 +2006,19 @@ def to_graphviz(
             dpi=dpi,
         )
 
+    def _repr_pretty_(self, p, cycle):
+        from pymc.printing import str_for_model
+
+        output = str_for_model(self)
+        # Find newlines and replace them with p.break_()
+        # (see IPython.lib.pretty._repr_pprint)
+        lines = output.splitlines()
+        with p.group():
+            for idx, output_line in enumerate(lines):
+                if idx:
+                    p.break_()
+                p.text(output_line)
+
 
 class BlockModelAccess(Model):
     """Can be used to prevent user access to Model contexts."""
@@ -2240,19 +2245,6 @@ def Deterministic(name, var, model=None, dims=None):
     var = var.copy(model.name_for(name))
     model.deterministics.append(var)
     model.add_named_variable(var, dims)
-
-    from pymc.printing import str_for_potential_or_deterministic
-
-    var.str_repr = types.MethodType(
-        functools.partial(str_for_potential_or_deterministic, dist_name="Deterministic"), var
-    )
-    var._repr_latex_ = types.MethodType(
-        functools.partial(
-            str_for_potential_or_deterministic, dist_name="Deterministic", formatting="latex"
-        ),
-        var,
-    )
-
     return var
 
 
@@ -2365,16 +2357,4 @@ def normal_logp(value, mu, sigma):
     model.potentials.append(var)
     model.add_named_variable(var, dims)
 
-    from pymc.printing import str_for_potential_or_deterministic
-
-    var.str_repr = types.MethodType(
-        functools.partial(str_for_potential_or_deterministic, dist_name="Potential"), var
-    )
-    var._repr_latex_ = types.MethodType(
-        functools.partial(
-            str_for_potential_or_deterministic, dist_name="Potential", formatting="latex"
-        ),
-        var,
-    )
-
     return var

pymc/printing.py

@@ -35,7 +35,10 @@
 
 
 def str_for_dist(
-    dist: TensorVariable, formatting: str = "plain", include_params: bool = True
+    dist: TensorVariable,
+    formatting: str = "plain",
+    include_params: bool = True,
+    model: Model | None = None,
 ) -> str:
     """Make a human-readable string representation of a Distribution in a model.
 
@@ -47,12 +50,12 @@ def str_for_dist(
             dist.owner.op, "extended_signature", None
         ):
             dist_args = [
-                _str_for_input_var(x, formatting=formatting)
+                _str_for_input_var(x, formatting=formatting, model=model)
                 for x in dist.owner.op.dist_params(dist.owner)
             ]
         else:
             dist_args = [
-                _str_for_input_var(x, formatting=formatting)
+                _str_for_input_var(x, formatting=formatting, model=model)
                 for x in dist.owner.inputs
                 if not isinstance(x.type, RandomType | NoneTypeT)
             ]
@@ -106,7 +109,7 @@ def str_for_model(model: Model, formatting: str = "plain", include_params: bool
     including parameter values.
     """
     # Wrap functions to avoid confusing typecheckers
-    sfd = partial(str_for_dist, formatting=formatting, include_params=include_params)
+    sfd = partial(str_for_dist, formatting=formatting, include_params=include_params, model=model)
     sfp = partial(
         str_for_potential_or_deterministic, formatting=formatting, include_params=include_params
     )
@@ -169,18 +172,14 @@ def str_for_potential_or_deterministic(
             return rf"{print_name} ~ {dist_name}"
 
 
-def _str_for_input_var(var: Variable, formatting: str) -> str:
+def _str_for_input_var(var: Variable, formatting: str, model: Model | None = None) -> str:
     # Avoid circular import
     from pymc.distributions.distribution import SymbolicRandomVariable
 
     def _is_potential_or_deterministic(var: Variable) -> bool:
-        if not hasattr(var, "str_repr"):
-            return False
-        try:
-            return var.str_repr.__func__.func is str_for_potential_or_deterministic
-        except AttributeError:
-            # in case other code overrides str_repr, fallback
+        if model is None:
             return False
+        return var in model.deterministics or var in model.potentials
 
     if isinstance(var, Constant | SharedVariable):
         return _str_for_constant(var, formatting)
@@ -190,18 +189,18 @@ def _is_potential_or_deterministic(var: Variable) -> bool:
         # show the names for RandomVariables, Deterministics, and Potentials, rather
         # than the full expression
         assert isinstance(var, TensorVariable)
-        return _str_for_input_rv(var, formatting)
+        return _str_for_input_rv(var, formatting, model=model)
     elif isinstance(var.owner.op, DimShuffle):
-        return _str_for_input_var(var.owner.inputs[0], formatting)
+        return _str_for_input_var(var.owner.inputs[0], formatting, model=model)
     else:
         return _str_for_expression(var, formatting)
 
 
-def _str_for_input_rv(var: TensorVariable, formatting: str) -> str:
+def _str_for_input_rv(var: TensorVariable, formatting: str, model: Model | None = None) -> str:
     _str = (
         var.name
         if var.name is not None
-        else str_for_dist(var, formatting=formatting, include_params=True)
+        else str_for_dist(var, formatting=formatting, include_params=True, model=model)
     )
     if "latex" in formatting:
         return _latex_text_format(_latex_escape(_str.strip("$")))
@@ -277,37 +276,6 @@ def _latex_escape(text: str) -> str:
     return text.replace("$", r"\$")
 
 
-def _default_repr_pretty(obj: TensorVariable | Model, p, cycle):
-    """Handy plug-in method to instruct IPython-like REPLs to use our str_repr above."""
-    # we know that our str_repr does not recurse, so we can ignore cycle
-    try:
-        if not hasattr(obj, "str_repr"):
-            raise AttributeError
-        output = obj.str_repr()
-        # Find newlines and replace them with p.break_()
-        # (see IPython.lib.pretty._repr_pprint)
-        lines = output.splitlines()
-        with p.group():
-            for idx, output_line in enumerate(lines):
-                if idx:
-                    p.break_()
-                p.text(output_line)
-    except AttributeError:
-        # the default fallback option (no str_repr method)
-        IPython.lib.pretty._repr_pprint(obj, p, cycle)
-
-
-try:
-    # register our custom pretty printer in ipython shells
-    import IPython
-
-    IPython.lib.pretty.for_type(TensorVariable, _default_repr_pretty)
-    IPython.lib.pretty.for_type(Model, _default_repr_pretty)
-except (ModuleNotFoundError, AttributeError):
-    # no ipython shell
-    pass
-
-
 def _format_underscore(variable: str) -> str:
     """Escapes all unescaped underscores in the variable name for LaTeX representation."""
     return re.sub(r"(?<!\\)_", r"\\_", variable)

tests/test_printing.py

@@ -30,15 +30,17 @@
 )
 from pymc.math import dot
 from pymc.model import Deterministic, Model, Potential
+from pymc.printing import str_for_dist, str_for_model
 from pymc.pytensorf import floatX
 
 
 class BaseTestStrAndLatexRepr:
     def test__repr_latex_(self):
-        for distribution, tex in zip(self.distributions, self.expected[("latex", True)]):
-            assert distribution._repr_latex_() == tex
+        for model_variable, tex in zip(self.distributions, self.expected[("latex", True)]):
+            if model_variable in self.model.basic_RVs:
+                assert str_for_dist(model_variable, formatting="latex", model=self.model) == tex
 
-        model_tex = self.model._repr_latex_()
+        model_tex = str_for_model(self.model, formatting="latex")
 
         # make sure each variable is in the model
         for tex in self.expected[("latex", True)]:
@@ -47,10 +49,11 @@ def test__repr_latex_(self):
 
     def test_str_repr(self):
         for str_format in self.formats:
-            for dist, text in zip(self.distributions, self.expected[str_format]):
-                assert dist.str_repr(*str_format) == text
+            for model_variable, text in zip(self.distributions, self.expected[str_format]):
+                if model_variable in self.model.basic_RVs:
+                    assert str_for_dist(model_variable, *str_format, model=self.model) == text
 
-            model_text = self.model.str_repr(*str_format)
+            model_text = str_for_model(self.model, *str_format)
             for text in self.expected[str_format]:
                 if str_format[0] == "latex":
                     for segment in text.strip("$").split(r"\sim"):
@@ -252,7 +255,7 @@ def test_model_latex_repr_three_levels_model():
             "censored_normal", normal_dist, lower=-2.0, upper=2.0, observed=[1, 0, 0.5]
         )
 
-    latex_repr = censored_model.str_repr(formatting="latex")
+    latex_repr = str_for_model(censored_model, formatting="latex")
     expected = [
         "$$",
         "\\begin{array}{rcl}",
@@ -270,7 +273,7 @@ def test_model_latex_repr_mixture_model():
         w = Dirichlet("w", [1, 1])
         mix = Mixture("mix", w=w, comp_dists=[Normal.dist(0.0, 5.0), StudentT.dist(7.0)])
 
-    latex_repr = mix_model.str_repr(formatting="latex")
+    latex_repr = str_for_model(mix_model, formatting="latex")
     expected = [
         "$$",
         "\\begin{array}{rcl}",
@@ -291,15 +294,15 @@ def test_model_repr_variables_without_monkey_patched_repr():
     model = Model()
     model.register_rv(x, "x")
 
-    str_repr = model.str_repr()
+    str_repr = str_for_model(model)
     assert str_repr == "x ~ Normal(0, 1)"
 
 
 def test_truncated_repr():
     with Model() as model:
         x = Truncated("x", Gamma.dist(1, 1), lower=0, upper=20)
 
-    str_repr = model.str_repr(include_params=False)
+    str_repr = str_for_model(model, include_params=False)
     assert str_repr == "x ~ TruncatedGamma"
 
 
@@ -315,7 +318,7 @@ def random(rng, mu, size):
         x = CustomDist("x", 0, dist=dist, class_name="CustomDistNormal")
         x = CustomDist("y", 0, random=random, class_name="CustomRandomNormal")
 
-    str_repr = model.str_repr(include_params=False)
+    str_repr = str_for_model(model, include_params=False)
     assert str_repr == "\n".join(["x ~ CustomDistNormal", "y ~ CustomRandomNormal"])
 
 
@@ -333,6 +336,6 @@ def test_latex_escaped_underscore(self):
         Ensures that all underscores in model variable names are properly escaped for LaTeX representation
         """
         model = self.simple_model()
-        model_str = model.str_repr(formatting="latex")
+        model_str = str_for_model(model, formatting="latex")
         assert "\\_" in model_str
         assert "_" not in model_str.replace("\\_", "")