SupportTreescope -> SupportsTreescope

src/evermore/effect.py

@@ -9,7 +9,7 @@
 
 from evermore.custom_types import OffsetAndScale
 from evermore.parameter import Parameter
-from evermore.visualization import SupportTreescope
+from evermore.visualization import SupportsTreescope
 
 __all__ = [
     "Effect",
@@ -24,7 +24,7 @@ def __dir__():
     return __all__
 
 
-class Effect(eqx.Module, SupportTreescope):
+class Effect(eqx.Module, SupportsTreescope):
     @abc.abstractmethod
     def __call__(self, parameter: PyTree[Parameter], hist: Array) -> OffsetAndScale: ...
 

src/evermore/modifier.py

@@ -15,7 +15,7 @@
 from evermore.effect import DEFAULT_EFFECT
 from evermore.parameter import Parameter
 from evermore.util import tree_stack
-from evermore.visualization import SupportTreescope
+from evermore.visualization import SupportsTreescope
 
 if TYPE_CHECKING:
     from evermore.effect import Effect
@@ -59,7 +59,7 @@ def __matmul__(self: ModifierLike, other: ModifierLike) -> Compose:
         return Compose(self, other)
 
 
-class ModifierBase(ApplyFn, MatMulCompose, AbstractModifier, SupportTreescope):
+class ModifierBase(ApplyFn, MatMulCompose, AbstractModifier, SupportsTreescope):
     """
     This serves as a base class for all modifiers.
     It automatically implements the __call__ method to apply the scale factors to the hist array

src/evermore/parameter.py

@@ -12,7 +12,7 @@
 from evermore.custom_types import PDFLike
 from evermore.pdf import Normal, Poisson
 from evermore.util import filter_tree_map
-from evermore.visualization import SupportTreescope
+from evermore.visualization import SupportsTreescope
 
 if TYPE_CHECKING:
     from evermore.modifier import Modifier
@@ -34,7 +34,7 @@ def __dir__():
     return __all__
 
 
-class Parameter(eqx.Module, SupportTreescope):
+class Parameter(eqx.Module, SupportsTreescope):
     """
     Implementation of a general Parameter class. The class is used to define the parameters of a statistical model.
     Key is the value attribute, which holds the actual value of the parameter. In additon,

src/evermore/pdf.py

@@ -8,7 +8,7 @@
 from jax.scipy.special import gammaln, xlogy
 from jaxtyping import Array, PRNGKeyArray
 
-from evermore.visualization import SupportTreescope
+from evermore.visualization import SupportsTreescope
 
 __all__ = [
     "PDF",
@@ -21,7 +21,7 @@ def __dir__():
     return __all__
 
 
-class PDF(eqx.Module, SupportTreescope):
+class PDF(eqx.Module, SupportsTreescope):
     @abstractmethod
     def log_prob(self, x: Array) -> Array: ...
 

src/evermore/staterror.py

@@ -14,7 +14,7 @@
 from evermore.parameter import NormalParameter, Parameter
 from evermore.pdf import Poisson
 from evermore.util import sum_over_leaves
-from evermore.visualization import SupportTreescope
+from evermore.visualization import SupportsTreescope
 
 __all__ = [
     "StatErrors",
@@ -25,7 +25,7 @@ def __dir__():
     return __all__
 
 
-class StatErrors(eqx.Module, SupportTreescope):
+class StatErrors(eqx.Module, SupportsTreescope):
     """
     Create staterror (barlow-beeston) parameters.
 

src/evermore/visualization.py

@@ -12,7 +12,7 @@
 )
 
 
-class SupportTreescope:
+class SupportsTreescope:
     def __treescope_repr__(
         self,
         path: str,
@@ -56,7 +56,6 @@ def handle_evermore_classes(
     # get colors for the background of the tree node
     def _treescope_color(node) -> str:
         """Returns the color of the tree node."""
-        from treescope import formatting_util
 
         type_string = type(node).__module__ + "." + type(node).__qualname__
         return formatting_util.color_from_string(type_string)