Introduce context manager to handle theme in notebook

vizro-core/src/vizro/_vizro.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-
+import plotly.io as pio
 import logging
 import warnings
 from pathlib import Path
@@ -12,6 +12,7 @@
 from vizro._constants import STATIC_URL_PREFIX
 from vizro.managers import data_manager, model_manager
 from vizro.models import Dashboard
+from vizro.models.types import _pio_templates_default
 
 logger = logging.getLogger(__name__)
 
@@ -84,6 +85,15 @@ def build(self, dashboard: Dashboard):
         if dashboard.title:
             self.dash.title = dashboard.title
 
+        # Set global template to vizro_light or vizro_dark.
+        # The choice between these is generally meaningless because chart colours in the two are identical, and
+        # everything else gets overridden in the post-fig creation layout.template update in Graph.__call__ and the
+        # clientside theme selector callback.
+        # Note this setting of global template isn't undone anywhere. If we really wanted to then we could try and
+        # put in some teardown code, but it would probably never be 100% reliable. Remember this template setting
+        # can't go in run() though since it's needed even in deployment.
+        pio.templates.default = dashboard.theme
+
         # Note that model instantiation and pre_build are independent of Dash.
         self._pre_build()
 

vizro-core/src/vizro/models/_components/graph.py

@@ -5,6 +5,8 @@
 from dash.exceptions import MissingCallbackContextException
 from plotly import graph_objects as go
 
+import plotly.io as pio
+
 try:
     from pydantic.v1 import Field, PrivateAttr, validator
 except ImportError:  # pragma: no cov
@@ -61,14 +63,26 @@ def __call__(self, **kwargs):
         if fig.layout.title.text is None:
             fig.update_layout(margin_t=24)
 
+        # Apply the template vizro_dark or vizro_light by setting fig.layout.template. This is exactly the same as
+        # what the clientside update_graph_theme callback does, and it would be nice if we could just use that by
+        # including Input(self.id, "figure") as input for that callback, but doing so leads to a small flicker between
+        # completion of this serverside callback and starting that clientside callbacks.
+        # Note that this does not fully set the template for plotly.express figures. Doing this post-fig creation update
+        # relies on the fact that we have already set pio.templates.default before the self.figure call
+        # above (but not with the _pio_templates_default context manager surrounding the above self.figure call,
+        # since that would alter global state).
+        # Possibly we should pass through the theme selector as an argument `template` in __call__ rather than fetching
+        # it from ctx here. Remember that passing it as self.figure(template) is not helpful though, because custom
+        # graph figures don't need a template argument, and the clientside them selector callback would override this
+        # anyway.
         # Possibly we should enforce that __call__ can only be used within the context of a callback, but it's easy
         # to just swallow up the error here as it doesn't cause any problems.
         try:
             # At the moment theme_selector is always present so this if statement is redundant, but possibly in
             # future we'll have callbacks that do Graph.__call__() without theme_selector set.
             if "theme_selector" in ctx.args_grouping.get("external", {}):
                 theme_selector_checked = ctx.args_grouping["external"]["theme_selector"]["value"]
-                fig.layout.template = themes.light if theme_selector_checked else themes.dark
+                fig.layout.template = "vizro_light" if theme_selector_checked else "vizro_dark"
         except MissingCallbackContextException:
             logger.info("fig.update_layout called outside of callback context.")
         return fig

vizro-core/src/vizro/models/_dashboard.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import plotly.io as pio
+
 import logging
 from functools import partial
 from pathlib import Path

vizro-core/src/vizro/models/types.py

@@ -6,6 +6,7 @@
 import functools
 import importlib
 import inspect
+from contextlib import contextmanager
 from datetime import date
 from typing import Any, Dict, List, Literal, Protocol, Union, runtime_checkable
 
@@ -250,6 +251,39 @@ def _check_type(cls, captured_callable: CapturedCallable, field: ModelField) ->
         return captured_callable
 
 
+@contextmanager
+def _pio_templates_default(default: Literal["vizro_light", "vizro_dark"]):
+    """Sets pio.templates.default and then reverts it.
+
+    This is to ensure that in a Jupyter notebook captured charts look the same as when they're in the dashboard. When
+    the context manager exits the global theme is reverted just to keep things clean (e.g. if you really wanted to,
+    you could compare a captured vs. non-captured chart in the same Python session).
+
+    This works even if users have tweaked the templates, so long as pio.templates has been updated correctly and you
+    refer to template by name rather than trying to take from vizro.themes.
+
+    If pio.templates.default has already been set to vizro_dark or vizro_light then no change is made to allow a user
+    to set these without it being overridden.
+    """
+
+    old_default = pio.templates.default
+    template_changed = False
+    # If the user has set pio.templates.default to a vizro theme already, no need to change it.
+    if old_default not in ["vizro_dark", "vizro_light"]:
+        template_changed = True
+        pio.templates.default = default
+
+    # Revert the template. This is done in a try/finally so that if the code wrapped inside the context manager (i.e.
+    # plotting functions) raises an exception, pio.templates.default is still reverted. This is not very important
+    # but easy to achieve.
+    try:
+        # This will always be vizro_light or vizro_dark and corresponds to the default theme that has been set.
+        yield pio.templates.default
+    finally:
+        if template_changed:
+            pio.templates.default = old_default
+
+
 class capture:
     """Captures a function call to create a [`CapturedCallable`][vizro.models.types.CapturedCallable].
 
@@ -334,23 +368,25 @@ def wrapped(*args, **kwargs) -> _DashboardReadyFigure:
                     fig = _DashboardReadyFigure()
                 else:
                     # Standard case for px.scatter(df: pd.DataFrame).
-                    fig = func(*args, **kwargs)
+                    # Set theme for the figure that gets shown in a Jupyter notebook. This is to ensure that in a
+                    # Jupyter notebook captured charts look the same as when they're in the dashboard. To mimic this,
+                    # we first use _pio_templates_default to set the global theme, as is done in the dashboard, and then
+                    # do the fig.layout.template update that is achieved by the theme selector.
+                    # We don't want to update the captured_callable in the same way, since it's only used inside the
+                    # dashboard, at which point the global pio.templates.default is always set anyway according to
+                    # the dashboard theme and then updated according to the theme selector.
+                    with _pio_templates_default("vizro_dark") as default_template:
+                        fig = func(*args, **kwargs)
+                    # Update the fig.layout.template just to ensure absolute consistency with how the dashboard
+                    # works.
+                    # The only exception here is the edge case that the user has specified template="vizro_light" or
+                    # "vizro_dark" in the plotting function, in which case we don't want to change it. This makes
+                    # it easier for a user to try out both themes simultaneously in a notebook.
+                    if fig.layout.template not in (pio.templates["vizro_dark"], pio.templates["vizro_light"]):
+                        fig.layout.template = default_template
                     fig.__class__ = _DashboardReadyFigure
 
                 fig._captured_callable = captured_callable
-
-                # For Jupyter notebook users, we want the chart to show as if it's in a dashboard, so apply the same
-                # theme that we do in the dashboard. We use the vizro_dark theme by default. The only exception to this
-                # is if the user has explicitly chosen to use vizro_light (could be through setting
-                # pio.default.templates or setting template="vizro_light" or any other way). In this case we don't
-                # want to change the template to vizro_dark.
-                # This works even if users have tweaked the templates, so long as pio.templates has been updated
-                # correctly and you refer to template by name rather than trying to take from vizro.themes.
-                # We don't want to update the captured_callable in the same way, since it's only used inside the
-                # dashboard, at which point the theme always gets set according to the theme selector.
-                if fig.layout.template != pio.templates["vizro_light"]:
-                    fig.layout.template = pio.templates["vizro_dark"]
-
                 return fig
 
             return wrapped