From 41b68da5696bc6b986c2e87c21bb06cff414dc7b Mon Sep 17 00:00:00 2001
From: Jon Mease <jonmmease@gmail.com>
Date: Tue, 2 Jan 2024 08:35:34 -0500
Subject: [PATCH] Add embed_options to JupyterChart and pass them through in
 "jupyter" renderer

---
 altair/jupyter/js/index.js          | 4 +++-
 altair/jupyter/jupyter_chart.py     | 8 +++++++-
 altair/vegalite/v5/display.py       | 8 ++++++--
 tests/vegalite/v5/test_renderers.py | 9 +++++++++
 4 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/altair/jupyter/js/index.js b/altair/jupyter/js/index.js
index 9202e99be..779514f86 100644
--- a/altair/jupyter/js/index.js
+++ b/altair/jupyter/js/index.js
@@ -30,10 +30,11 @@ export async function render({ model, el }) {
             model.save_changes();
             return;
         }
+        let embedOptions = structuredClone(model.get("embed_options")) ?? undefined;
 
         let api;
         try {
-            api = await vegaEmbed(el, spec);
+            api = await vegaEmbed(el, spec, embedOptions);
         } catch (error) {
             showError(error)
             return;
@@ -134,6 +135,7 @@ export async function render({ model, el }) {
     }
 
     model.on('change:spec', reembed);
+    model.on('change:embed_options', reembed);
     model.on('change:debounce_wait', reembed);
     model.on('change:max_wait', reembed);
     await reembed();
diff --git a/altair/jupyter/jupyter_chart.py b/altair/jupyter/jupyter_chart.py
index fc973a5a8..cb5b8e0f6 100644
--- a/altair/jupyter/jupyter_chart.py
+++ b/altair/jupyter/jupyter_chart.py
@@ -2,7 +2,7 @@
 import anywidget
 import traitlets
 import pathlib
-from typing import Any, Set
+from typing import Any, Set, Optional
 
 import altair as alt
 from altair.utils._vegafusion_data import (
@@ -109,6 +109,7 @@ class JupyterChart(anywidget.AnyWidget):
     max_wait = traitlets.Bool(default_value=True).tag(sync=True)
     local_tz = traitlets.Unicode(default_value=None, allow_none=True).tag(sync=True)
     debug = traitlets.Bool(default_value=False)
+    embed_options = traitlets.Dict(default_value=None, allow_none=True).tag(sync=True)
 
     # Internal selection traitlets
     _selection_types = traitlets.Dict()
@@ -129,6 +130,7 @@ def __init__(
         debounce_wait: int = 10,
         max_wait: bool = True,
         debug: bool = False,
+        embed_options: Optional[dict] = None,
         **kwargs: Any,
     ):
         """
@@ -148,6 +150,9 @@ def __init__(
              sent until chart interactions have completed.
         debug: bool
              If True, debug messages will be printed
+        embed_options: dict
+             Options to pass to vega-embed.
+             See https://github.com/vega/vega-embed?tab=readme-ov-file#options
         """
         self.params = Params({})
         self.selections = Selections({})
@@ -156,6 +161,7 @@ def __init__(
             debounce_wait=debounce_wait,
             max_wait=max_wait,
             debug=debug,
+            embed_options=embed_options,
             **kwargs,
         )
 
diff --git a/altair/vegalite/v5/display.py b/altair/vegalite/v5/display.py
index b13d62e06..0b0a5b66a 100644
--- a/altair/vegalite/v5/display.py
+++ b/altair/vegalite/v5/display.py
@@ -86,13 +86,17 @@ def svg_renderer(spec: dict, **metadata) -> Dict[str, str]:
     )
 
 
-def jupyter_renderer(spec: dict):
+def jupyter_renderer(spec: dict, **metadata):
     """Render chart using the JupyterChart Jupyter Widget"""
     from altair import Chart, JupyterChart
 
+    embed_options = metadata.get("embed_options", None)
+
     # Need to ignore attr-defined mypy rule because mypy doesn't see _repr_mimebundle_
     # conditionally defined in AnyWidget
-    return JupyterChart(chart=Chart.from_dict(spec))._repr_mimebundle_()  # type: ignore[attr-defined]
+    return JupyterChart(
+        chart=Chart.from_dict(spec), embed_options=embed_options
+    )._repr_mimebundle_()  # type: ignore[attr-defined]
 
 
 html_renderer = HTMLRenderer(
diff --git a/tests/vegalite/v5/test_renderers.py b/tests/vegalite/v5/test_renderers.py
index 8a62829e1..998a958c9 100644
--- a/tests/vegalite/v5/test_renderers.py
+++ b/tests/vegalite/v5/test_renderers.py
@@ -63,3 +63,12 @@ def test_json_renderer_embed_options(chart, renderer="json"):
             assert set(bundle.keys()) == {mimetype, "text/plain"}
             assert bundle[mimetype] == spec
             assert metadata == {mimetype: {"option": "foo"}}
+
+
+def test_jupyter_renderer_mimetype(chart, renderer="jupyter"):
+    """Test that we get the expected widget mimetype when the jupyter renderer is enabled"""
+    with alt.renderers.enable(renderer):
+        assert (
+            "application/vnd.jupyter.widget-view+json"
+            in chart._repr_mimebundle_(None, None)[0]
+        )