From 54191eb6f7a3df1353a3f8029f93bce49f2050c7 Mon Sep 17 00:00:00 2001 From: Jake VanderPlas Date: Sun, 15 Mar 2020 08:28:36 -0700 Subject: [PATCH] ENH: support embed_options in html & selenium savers --- altair_saver/savers/_html.py | 16 +++++++------- altair_saver/savers/_node.py | 4 ++++ altair_saver/savers/_saver.py | 10 ++++++++- altair_saver/savers/_selenium.py | 30 ++++++++++++-------------- altair_saver/savers/tests/test_html.py | 10 ++++++--- 5 files changed, 42 insertions(+), 28 deletions(-) diff --git a/altair_saver/savers/_html.py b/altair_saver/savers/_html.py index 78d4762..4a82dd5 100644 --- a/altair_saver/savers/_html.py +++ b/altair_saver/savers/_html.py @@ -18,7 +18,7 @@
@@ -42,7 +42,7 @@
@@ -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"), diff --git a/altair_saver/savers/_node.py b/altair_saver/savers/_node.py index 2ad7bcf..eab0d0a 100644 --- a/altair_saver/savers/_node.py +++ b/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'") diff --git a/altair_saver/savers/_saver.py b/altair_saver/savers/_saver.py index fc89334..5f467d0 100644 --- a/altair_saver/savers/_saver.py +++ b/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: diff --git a/altair_saver/savers/_selenium.py b/altair_saver/savers/_selenium.py index c8fe651..4a1680e 100644 --- a/altair_saver/savers/_selenium.py +++ b/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,11 +155,11 @@ 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: @@ -168,12 +167,11 @@ def __init__( 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"] diff --git a/altair_saver/savers/tests/test_html.py b/altair_saver/savers/tests/test_html.py index c74faf6..6c4c7cd 100644 --- a/altair_saver/savers/tests/test_html.py +++ b/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("") assert json.dumps(data["vega-lite"]) in html + assert f"const embedOpt = {json.dumps(embed_options or {})}" in html def test_bad_format() -> None: