Merge pull request #554 from jakevdp/selenium

Add programmatic export utility based on selenium and chrome-headless
vega · Mar 5, 2018 · f65cf16 · f65cf16
2 parents cf4ebb6 + 5b78267
commit f65cf16
Show file tree

Hide file tree

Showing 4 changed files with 150 additions and 1 deletion.
diff --git a/altair/utils/__init__.py b/altair/utils/__init__.py
@@ -2,4 +2,6 @@
 
 from .codegen import CodeGen
 
-from .plugin_registry import PluginRegistry
+from .plugin_registry import PluginRegistry
+
+from .headless import save_spec
diff --git a/altair/utils/headless.py b/altair/utils/headless.py
@@ -0,0 +1,101 @@
+"""
+Utilities that use selenium + chrome headless to save figures
+"""
+
+import base64
+import io
+import json
+import os
+import tempfile
+
+from .importing import attempt_import
+
+
+HTML_TEMPLATE = """
+<!DOCTYPE html>
+<html>
+<head>
+  <title>Embedding Vega-Lite</title>
+  <script src="https://cdn.jsdelivr.net/npm/vega@3.0.10"></script>
+  <script src="https://cdn.jsdelivr.net/npm/vega-lite@2.1.3"></script>
+  <script src="https://cdn.jsdelivr.net/npm/vega-embed@3.0.0"></script>
+</head>
+<body>
+  <div id="vis"></div>
+</body>
+</html>
+"""
+
+EMBED_CODE = """
+vegaEmbed("#vis", {spec}).then(function(result) {{
+    window.view = result.view;
+}}).catch(console.error);
+"""
+
+CONVERT_CODE = """
+window.view.toCanvas().then(function(canvas) {
+    window.png_render = canvas.toDataURL('image/png');
+})
+"""
+
+EXTRACT_CODE = """
+return window.png_render;
+"""
+
+def save_spec(spec, filename, format=None, mode='vega-lite'):
+    """Save a spec to file
+
+    Parameters
+    ----------
+    spec : dict
+        a dictionary representing a vega-lite plot spec
+    filename : string
+        the filename at which the result will be saved
+    format : string (optional)
+        the file format to be saved. If not specified, it will be inferred
+        from the extension of filename
+    mode : string
+        Whether the spec is 'vega' or 'vega-lite'.
+        Currently only mode='vega-lite' is supported.
+
+    Note
+    ----
+    This requires the pillow, selenium, and chrome headless packages to be
+    installed.
+    """
+    # TODO: remove PIL dependency?
+    # TODO: use SVG renderer when it makes sense
+    # TODO: support mode='vega'
+    # TODO: allow package versions to be specified
+    # TODO: detect local Jupyter caches of JS packages?
+
+    Image = attempt_import('PIL.Image',
+                           'save_spec requires the pillow package')
+    webdriver = attempt_import('selenium.webdriver',
+                               'save_spec requires the selenium package')
+    Options = attempt_import('selenium.webdriver.chrome.options',
+                             'save_spec requires the selenium package').Options
+
+    try:
+        chrome_options = Options()
+        chrome_options.add_argument("--headless")
+        driver = webdriver.Chrome(chrome_options=chrome_options)
+        try:
+            fd, name = tempfile.mkstemp(suffix='.html', text=True)
+            with open(name, 'w') as f:
+                f.write(HTML_TEMPLATE)
+            driver.get("file://" + name)
+            driver.execute_script(EMBED_CODE.format(spec=json.dumps(spec)))
+            driver.execute_script(CONVERT_CODE)
+            png_base64 = driver.execute_script(EXTRACT_CODE)
+        finally:
+            os.remove(name)
+    finally:
+        driver.close()
+
+    out = io.BytesIO()
+    metadata, image = png_base64.split(',')
+    base64.decode(io.BytesIO(image.encode()), out)
+    out.seek(0)
+    image = Image.open(out)
+    image.save(filename)
diff --git a/altair/utils/importing.py b/altair/utils/importing.py
@@ -0,0 +1,31 @@
+"""
+Tools for importing dependencies
+"""
+import importlib
+
+
+def attempt_import(module, error_message):
+    """Attempt import of a dependency
+
+    If the dependency is not available, raise a RuntimeError with the appropriate message
+
+    Parameters
+    ----------
+    module : string
+        the module name
+    error_message : string
+        the error to raise if the module is not importable
+
+    Returns
+    -------
+    mod : ModuleType
+        the imported module
+
+    Raises
+    ------
+    RuntimeError
+    """
+    try:
+        return importlib.import_module(module)
+    except ImportError:
+        raise RuntimeError(error_message)
diff --git a/altair/utils/tests/test_importing.py b/altair/utils/tests/test_importing.py
@@ -0,0 +1,15 @@
+import pytest
+
+import pandas as pd
+
+from ..importing import attempt_import
+
+
+def test_import_module():
+    pandas = attempt_import('pandas', "Pandas package is not installed")
+    assert pandas is pd
+
+    with pytest.raises(RuntimeError) as err:
+        attempt_import('a_module_which_should_not_exist',
+                       'this module does not exist')
+    assert(str(err.value) == 'this module does not exist')