Merge pull request #30 from jakevdp/embed-opt

ENH: support embed_options in html & selenium savers

altair_saver/savers/_html.py

@@ -18,7 +18,7 @@
 <div id="vis"></div>
 <script type="text/javascript">
   const spec = {spec};
-  const embedOpt = {embed_opt};
+  const embedOpt = {embed_options};
   vegaEmbed('#vis', spec, embedOpt).catch(console.error);
 </script>
 </body>
@@ -42,7 +42,7 @@
 <div id="vis"></div>
 <script type="text/javascript">
   const spec = {spec};
-  const embedOpt = {embed_opt};
+  const embedOpt = {embed_options};
   vegaEmbed('#vis', spec, embedOpt).catch(console.error);
 </script>
 </body>
@@ -58,26 +58,26 @@ class HTMLSaver(Saver):
     valid_formats: List[str] = ["html"]
     _package_versions: Dict[str, str]
     _inline: bool
-    _embed_opt: JSONDict
+    _embed_options: JSONDict
 
     def __init__(
         self,
         spec: JSONDict,
         mode: Optional[str] = None,
+        embed_options: Optional[JSONDict] = None,
         inline: bool = False,
-        embed_opt: Optional[JSONDict] = None,
         vega_version: str = alt.VEGA_VERSION,
         vegalite_version: str = alt.VEGALITE_VERSION,
         vegaembed_version: str = alt.VEGAEMBED_VERSION,
     ) -> None:
         self._inline = inline
-        self._embed_opt = embed_opt or {}
+        self._embed_options = embed_options or {}
         self._package_versions = {
             "vega": vega_version,
             "vega-lite": vegalite_version,
             "vega-embed": vegaembed_version,
         }
-        super().__init__(spec=spec, mode=mode)
+        super().__init__(spec=spec, mode=mode, embed_options=embed_options)
 
     def _package_url(self, package: str) -> str:
         return CDN_URL.format(package=package, version=self._package_versions[package])
@@ -90,7 +90,7 @@ def _mimebundle(self, fmt: str) -> Mimebundle:
         if self._inline:
             html = INLINE_HTML_TEMPLATE.format(
                 spec=json.dumps(self._spec),
-                embed_opt=json.dumps(self._embed_opt),
+                embed_options=json.dumps(self._embed_options),
                 vega_version=self._package_versions["vega"],
                 vegalite_version=self._package_versions["vega-lite"],
                 vegaembed_version=self._package_versions["vega-embed"],
@@ -105,7 +105,7 @@ def _mimebundle(self, fmt: str) -> Mimebundle:
         else:
             html = HTML_TEMPLATE.format(
                 spec=json.dumps(self._spec),
-                embed_opt=json.dumps(self._embed_opt),
+                embed_options=json.dumps(self._embed_options),
                 vega_url=self._package_url("vega"),
                 vegalite_url=self._package_url("vega-lite"),
                 vegaembed_url=self._package_url("vega-embed"),

altair_saver/savers/_node.py

@@ -2,6 +2,7 @@
 import json
 import shutil
 from typing import List
+import warnings
 
 from altair_saver.savers import Saver
 from altair_saver._utils import (
@@ -79,6 +80,9 @@ def enabled(cls) -> bool:
 
     def _mimebundle(self, fmt: str) -> Mimebundle:
         """Return a mimebundle with a single mimetype."""
+        if self._embed_options:
+            warnings.warn("embed_options are not supported for method='node'.")
+
         if self._mode not in ["vega", "vega-lite"]:
             raise ValueError("mode must be either 'vega' or 'vega-lite'")
 

altair_saver/savers/_saver.py

@@ -23,8 +23,15 @@ class Saver(metaclass=abc.ABCMeta):
     valid_formats: List[str] = []
     _spec: JSONDict
     _mode: str
+    _embed_options: JSONDict
 
-    def __init__(self, spec: JSONDict, mode: Optional[str] = None, **kwargs: Any):
+    def __init__(
+        self,
+        spec: JSONDict,
+        mode: Optional[str] = None,
+        embed_options: Optional[JSONDict] = None,
+        **kwargs: Any,
+    ):
         if kwargs:
             raise ValueError(f"Unhandled keyword arguments: {list(kwargs.keys())}")
         if mode is None:
@@ -33,6 +40,7 @@ def __init__(self, spec: JSONDict, mode: Optional[str] = None, **kwargs: Any):
             raise ValueError("mode must be either 'vega' or 'vega-lite'")
         self._spec = spec
         self._mode = mode
+        self._embed_options = embed_options or {}
 
     @abc.abstractmethod
     def _mimebundle(self, fmt: str) -> Mimebundle:

altair_saver/savers/_selenium.py

@@ -38,14 +38,13 @@ class JavascriptError(RuntimeError):
 """
 
 EXTRACT_CODE = """
-var spec = arguments[0];
-var mode = arguments[1];
-var scaleFactor = arguments[2];
-var format = arguments[3];
-var done = arguments[4];
+let spec = arguments[0];
+const embedOpt = arguments[1];
+const format = arguments[2];
+const done = arguments[3];
 
 if (format === 'vega') {
-    if (mode === 'vega-lite') {
+    if (embedOpt['mode'] === 'vega-lite') {
         vegaLite = (typeof vegaLite === "undefined") ? vl : vegaLite;
         try {
             const compiled = vegaLite.compile(spec);
@@ -57,10 +56,10 @@ class JavascriptError(RuntimeError):
     done({result: spec});
 }
 
-vegaEmbed('#vis', spec, {mode}).then(function(result) {
+vegaEmbed('#vis', spec, embedOpt).then(function(result) {
     if (format === 'png') {
         result.view
-            .toCanvas(scaleFactor)
+            .toCanvas()
             .then(function(canvas){return canvas.toDataURL('image/png');})
             .then(result => done({result}))
             .catch(function(err) {
@@ -69,14 +68,14 @@ class JavascriptError(RuntimeError):
             });
     } else if (format === 'svg') {
         result.view
-            .toSVG(scaleFactor)
+            .toSVG()
             .then(result => done({result}))
             .catch(function(err) {
                 console.error(err);
                 done({error: err.toString()});
             });
     } else {
-        error = "Unrecognized format: " + format;
+        const error = "Unrecognized format: " + format;
         console.error(error);
         done({error});
     }
@@ -156,24 +155,23 @@ def __init__(
         self,
         spec: JSONDict,
         mode: Optional[str] = None,
+        embed_options: Optional[JSONDict] = None,
         vega_version: str = alt.VEGA_VERSION,
         vegalite_version: str = alt.VEGALITE_VERSION,
         vegaembed_version: str = alt.VEGAEMBED_VERSION,
         driver_timeout: int = 20,
-        scale_factor: float = 1,
         webdriver: Optional[Union[str, WebDriver]] = None,
         offline: bool = True,
     ) -> None:
         self._vega_version = vega_version
         self._vegalite_version = vegalite_version
         self._vegaembed_version = vegaembed_version
         self._driver_timeout = driver_timeout
-        self._scale_factor = scale_factor
         self._webdriver = (
             self._select_webdriver(driver_timeout) if webdriver is None else webdriver
         )
         self._offline = offline
-        super().__init__(spec=spec, mode=mode)
+        super().__init__(spec=spec, mode=mode, embed_options=embed_options)
 
     @classmethod
     def _select_webdriver(cls, driver_timeout: int) -> Optional[str]:
@@ -256,9 +254,9 @@ def _extract(self, fmt: str) -> MimeType:
                 raise RuntimeError(
                     f"Internet connection required for saving chart as {fmt} with offline=False."
                 )
-        result = driver.execute_async_script(
-            EXTRACT_CODE, self._spec, self._mode, self._scale_factor, fmt
-        )
+        opt = self._embed_options.copy()
+        opt["mode"] = self._mode
+        result = driver.execute_async_script(EXTRACT_CODE, self._spec, opt, fmt)
         if "error" in result:
             raise JavascriptError(result["error"])
         return result["result"]

altair_saver/savers/tests/test_html.py

@@ -1,7 +1,7 @@
 import io
 import json
 import os
-from typing import Any, Dict, IO, Iterator, Tuple
+from typing import Any, Dict, IO, Iterator, Optional, Tuple
 
 from altair_data_server import Provider
 from PIL import Image
@@ -49,14 +49,18 @@ def get_testcases() -> Iterator[Tuple[str, Dict[str, Any]]]:
 
 
 @pytest.mark.parametrize("inline", [True, False])
+@pytest.mark.parametrize("embed_options", [None, {"theme": "dark"}])
 @pytest.mark.parametrize("case, data", get_testcases())
-def test_html_saver(case: str, data: Dict[str, Any], inline: bool) -> None:
-    saver = HTMLSaver(data["vega-lite"], inline=inline)
+def test_html_saver(
+    case: str, data: Dict[str, Any], embed_options: Optional[dict], inline: bool
+) -> None:
+    saver = HTMLSaver(data["vega-lite"], inline=inline, embed_options=embed_options)
     bundle = saver.mimebundle("html")
     html = bundle.popitem()[1]
     assert isinstance(html, str)
     assert html.strip().startswith("<!DOCTYPE html>")
     assert json.dumps(data["vega-lite"]) in html
+    assert f"const embedOpt = {json.dumps(embed_options or {})}" in html
 
 
 def test_bad_format() -> None: