Add a new template parameter to ridgeplot()

Allow users to specify a Plotly figure template (by name, as a
go.layout.Template object, or as a dict) directly through the new
`template` parameter. When no explicit `colorscale` is specified, the
default colorscale is now also inferred from the specified template.

Author: tpvasconcelos

src/ridgeplot/_color/colorscale.py

@@ -14,6 +14,8 @@
 if TYPE_CHECKING:
     from collections.abc import Collection
 
+    import plotly.graph_objects as go
+
     from ridgeplot._types import Color, ColorScale
 
 
@@ -43,19 +45,29 @@ def validate_coerce(self, v: Any) -> ColorScale:
         return cast("ColorScale", coerced)
 
 
-def infer_default_colorscale() -> ColorScale | Collection[Color] | str:
+def infer_default_colorscale(
+    template: go.layout.Template | None = None,
+) -> ColorScale | Collection[Color] | str:
+    if template is None:
+        template = default_plotly_template()
     return validate_coerce_colorscale(
-        default_plotly_template().layout.colorscale.sequential or px.colors.sequential.Viridis
+        template.layout.colorscale.sequential or px.colors.sequential.Viridis
     )
 
 
 def validate_coerce_colorscale(
     colorscale: ColorScale | Collection[Color] | str | None,
+    template: go.layout.Template | None = None,
 ) -> ColorScale:
     """Convert mixed colorscale representations to the canonical
-    :data:`ColorScale` format."""
+    :data:`ColorScale` format.
+
+    If ``colorscale`` is None, the default colorscale is inferred from the
+    given ``template`` (or from the current default Plotly template, if no
+    template is specified).
+    """
     if colorscale is None:
-        colorscale = infer_default_colorscale()
+        colorscale = infer_default_colorscale(template=template)
     return ColorscaleValidator().validate_coerce(colorscale)
 
 

src/ridgeplot/_figure_factory.py

@@ -14,6 +14,7 @@
 )
 from ridgeplot._obj.traces import get_trace_cls
 from ridgeplot._obj.traces.base import ColoringContext
+from ridgeplot._template import validate_coerce_template
 from ridgeplot._types import (
     Color,
     ColorScale,
@@ -37,7 +38,7 @@
 if TYPE_CHECKING:
     from collections.abc import Collection
 
-    from ridgeplot._types import Densities
+    from ridgeplot._types import Densities, PlotlyTemplate
 
 
 def normalise_trace_types(
@@ -80,8 +81,11 @@ def _update_layout(
     xpad: float,
     x_max: float,
     x_min: float,
+    template: go.layout.Template | None,
 ) -> go.Figure:
     """Update figure's layout."""
+    if template is not None:
+        fig.update_layout(template=template)
     fig.update_layout(
         legend=dict(traceorder="normal"),
     )
@@ -130,6 +134,7 @@ def create_ridgeplot(
     line_width: float | None,
     spacing: float,
     xpad: float,
+    template: PlotlyTemplate | None,
 ) -> go.Figure:
     # ==============================================================
     # ---  Get clean and validated input arguments
@@ -169,7 +174,8 @@ def create_ridgeplot(
     line_width = float(line_width) if line_width is not None else None
     spacing = float(spacing)
     xpad = float(xpad)
-    colorscale = validate_coerce_colorscale(colorscale)
+    template = validate_coerce_template(template)
+    colorscale = validate_coerce_colorscale(colorscale, template=template)
 
     # ==============================================================
     # ---  Build the figure
@@ -233,5 +239,6 @@ def create_ridgeplot(
         xpad=xpad,
         x_max=x_max,
         x_min=x_min,
+        template=template,
     )
     return fig

src/ridgeplot/_ridgeplot.py

@@ -31,6 +31,7 @@
         ColorScale,
         LabelsArray,
         NormalisationOption,
+        PlotlyTemplate,
         SampleWeights,
         SampleWeightsArray,
         ShallowDensities,
@@ -120,6 +121,7 @@ def ridgeplot(
     line_width: float | None = None,
     spacing: float = 0.5,
     xpad: float = 0.05,
+    template: PlotlyTemplate | None = None,
     # Deprecated parameters
     coloralpha: float | None | MissingType = MISSING,
     linewidth: float | MissingType = MISSING,
@@ -140,6 +142,8 @@ def ridgeplot(
         https://www.statsmodels.org/stable/_modules/statsmodels/nonparametric/bandwidths.html
     .. _Plotly's built-in color-scales:
         https://plotly.com/python/builtin-colorscales/
+    .. _Plotly's theming and templates guide:
+        https://plotly.com/python/templates/
     .. _ragged:
        https://en.wikipedia.org/wiki/Jagged_array
 
@@ -290,8 +294,8 @@ def ridgeplot(
         ``["rgb(255, 0, 0)", "blue", "hsl(120, 100%, 50%)"]``). The list will
         ultimately be converted into a :data:`~ridgeplot._types.ColorScale`
         object, assuming the colors provided are evenly spaced. If not specified
-        (default), the color scale will be inferred from current Plotly
-        template.
+        (default), the color scale will be inferred from the current Plotly
+        template (see the :paramref:`.template` parameter).
 
     colormode : "fillgradient" or SolidColormode
         This parameter controls the logic used for the coloring of each
@@ -395,6 +399,20 @@ def ridgeplot(
         units of the range between the minimum and maximum x-values from all
         distributions.
 
+    template : PlotlyTemplate or None
+        The Plotly template to use in this figure. This can be the name of a
+        registered template (e.g., ``"plotly_dark"``), a
+        :class:`plotly.graph_objects.layout.Template
+        <plotly.graph_objs.layout.Template>` object, or a dictionary with a
+        template's properties. See `Plotly's theming and templates guide`_
+        for more details. If not specified (default), the plot will be
+        rendered using Plotly's current default template (i.e.,
+        ``plotly.io.templates.default``). Note that, unless a custom
+        :paramref:`.colorscale` is specified, the default colorscale will
+        also be inferred from this template.
+
+        .. versionadded:: 0.7.0
+
     coloralpha : float
 
         .. deprecated:: 0.2.0
@@ -511,5 +529,6 @@ def ridgeplot(
         line_width=line_width,
         spacing=spacing,
         xpad=xpad,
+        template=template,
     )
     return fig

src/ridgeplot/_template.py

@@ -0,0 +1,31 @@
+"""Plotly figure template utilities."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
+
+from _plotly_utils.basevalidators import BaseTemplateValidator
+
+if TYPE_CHECKING:
+    import plotly.graph_objects as go
+
+    from ridgeplot._types import PlotlyTemplate
+
+
+def validate_coerce_template(template: PlotlyTemplate | None) -> go.layout.Template | None:
+    """Convert mixed template representations into a
+    :class:`plotly.graph_objects.layout.Template
+    <plotly.graph_objs.layout.Template>` object.
+
+    ``None`` is passed through as-is, meaning that no template has been
+    specified and that Plotly's current default template should be used.
+    """
+    if template is None:
+        return None
+    validator = BaseTemplateValidator(
+        plotly_name="template",
+        parent_name="layout",
+        data_class_str="Template",
+        data_docs="",
+    )
+    return cast("go.layout.Template", validator.validate_coerce(template))

src/ridgeplot/_types.py

@@ -7,6 +7,7 @@
 from typing import Literal, TypeAlias
 
 import numpy as np
+import plotly.graph_objects as go
 from typing_extensions import Any, TypeIs, TypeVar
 
 # Snippet used to generate and store the image artifacts:
@@ -70,6 +71,16 @@
 available for the ridgeplot. See :paramref:`ridgeplot.ridgeplot.norm` for more
 details."""
 
+PlotlyTemplate: TypeAlias = go.layout.Template | dict[str, Any] | str
+"""A Plotly figure template can be represented by a
+:class:`plotly.graph_objects.layout.Template <plotly.graph_objs.layout.Template>`
+object, a dictionary with a template's properties, or the name of a registered
+template (e.g., ``"plotly_dark"``).
+
+See `Plotly's theming and templates guide
+<https://plotly.com/python/templates/>`_ for more details.
+"""
+
 # ========================================================
 # ---  Base nested Collection types (ragged arrays)
 # ========================================================

tests/unit/color/test_colorscale.py

@@ -3,6 +3,9 @@
 import re
 from typing import TYPE_CHECKING
 
+import plotly.express as px
+import plotly.graph_objects as go
+import plotly.io as pio
 import pytest
 
 from ridgeplot._color.colorscale import (
@@ -26,6 +29,20 @@ def test_infer_default_colorscale() -> None:
     assert infer_default_colorscale() == validate_coerce_colorscale("plasma")
 
 
+def test_infer_default_colorscale_from_template() -> None:
+    template = pio.templates["ggplot2"]
+    assert infer_default_colorscale(template=template) == validate_coerce_colorscale(
+        template.layout.colorscale.sequential
+    )
+
+
+def test_infer_default_colorscale_fallback_to_viridis() -> None:
+    template = go.layout.Template()
+    assert infer_default_colorscale(template=template) == validate_coerce_colorscale(
+        px.colors.sequential.Viridis
+    )
+
+
 # ==============================================================
 # ---  validate_coerce_colorscale()
 # ==============================================================

tests/unit/test_figure_factory.py

@@ -134,4 +134,5 @@ def test_densities_must_be_4d(self, densities: Densities) -> None:
                 line_width=...,  # type: ignore[reportArgumentType]
                 spacing=...,  # type: ignore[reportArgumentType]
                 xpad=...,  # type: ignore[reportArgumentType]
+                template=...,  # type: ignore[reportArgumentType]
             )

tests/unit/test_ridgeplot.py

@@ -4,6 +4,7 @@
 from typing import TYPE_CHECKING, Any
 
 import plotly.express as px
+import plotly.io as pio
 import pytest
 
 from ridgeplot import ridgeplot
@@ -171,6 +172,65 @@ def test_colorscale_invalid(invalid_colorscale: ColorScale | Collection[Color] |
         ridgeplot(samples=[[[1, 2, 3], [4, 5, 6]]], colorscale=invalid_colorscale)
 
 
+# ==============================================================
+# ---  param: template
+# ==============================================================
+
+
+def test_template_is_applied_to_the_figure() -> None:
+    fig = ridgeplot(samples=[[[1, 2, 3], [4, 5, 6]]], template="plotly_dark")
+    assert fig.layout.template == pio.templates["plotly_dark"]
+
+
+def test_template_equivalent_representations() -> None:
+    samples = [[[1, 2, 3], [4, 5, 6]]]
+    template = pio.templates["seaborn"]
+    assert (
+        ridgeplot(samples=samples, template="seaborn") ==
+        ridgeplot(samples=samples, template=template) ==
+        ridgeplot(samples=samples, template=template.to_plotly_json())
+    )  # fmt: skip
+
+
+def test_template_default() -> None:
+    # Not specifying a template should be equivalent to
+    # specifying Plotly's current default template
+    samples = [[[1, 2, 3], [4, 5, 6]]]
+    assert (
+        ridgeplot(samples=samples) ==
+        ridgeplot(samples=samples, template=pio.templates.default or "plotly")
+    )  # fmt: skip
+
+
+def test_template_sets_default_colorscale() -> None:
+    # The default colorscale should be inferred from the
+    # template specified via the `template` argument
+    samples = [[[1, 2, 3], [4, 5, 6]]]
+    template = pio.templates["ggplot2"]
+    assert (
+        ridgeplot(samples=samples, template="ggplot2") ==
+        ridgeplot(
+            samples=samples,
+            template="ggplot2",
+            colorscale=template.layout.colorscale.sequential,
+        )
+    )  # fmt: skip
+
+
+def test_template_explicit_colorscale_takes_precedence() -> None:
+    fig = ridgeplot(
+        samples=[[[1, 2, 3], [4, 5, 6]]],
+        template="ggplot2",
+        colorscale=(
+            (0.0, "rgb(10, 10, 10)"),
+            (1.0, "rgb(20, 20, 20)"),
+        ),
+        colormode="trace-index",
+    )
+    assert fig.data[1].fillcolor == "rgb(20, 20, 20)"
+    assert fig.data[3].fillcolor == "rgb(10, 10, 10)"
+
+
 # ==============================================================
 # ---  param: color_discrete_map
 # ==============================================================

tests/unit/test_template.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+import plotly.graph_objects as go
+import plotly.io as pio
+import pytest
+
+from ridgeplot._template import validate_coerce_template
+
+
+def test_none_passthrough() -> None:
+    assert validate_coerce_template(None) is None
+
+
+def test_registered_template_name() -> None:
+    assert validate_coerce_template("plotly_dark") == pio.templates["plotly_dark"]
+
+
+# Merging templates causes Plotly to internally instantiate all
+# registered trace types, including the deprecated Scattermapbox
+@pytest.mark.filterwarnings(r"ignore:\*scattermapbox\* is deprecated.*:DeprecationWarning")
+def test_merged_template_names() -> None:
+    assert validate_coerce_template("plotly_white+ggplot2") == pio.templates["plotly_white+ggplot2"]
+
+
+def test_template_object() -> None:
+    template = pio.templates["seaborn"]
+    assert validate_coerce_template(template) == template
+
+
+def test_template_dict() -> None:
+    template = pio.templates["seaborn"]
+    coerced = validate_coerce_template(template.to_plotly_json())
+    assert isinstance(coerced, go.layout.Template)
+    assert coerced == template
+
+
+def test_invalid_template_name() -> None:
+    with pytest.raises(ValueError, match=r"Invalid value .* received for the 'template' property"):
+        validate_coerce_template("this-is-not-a-registered-template")
+
+
+def test_invalid_template_type() -> None:
+    with pytest.raises(ValueError, match=r"Invalid value .* received for the 'template' property"):
+        validate_coerce_template(42)  # pyright: ignore[reportArgumentType]