Repository: vega/altair

Issue:
Move chart saving to an extension
Currently, saving charts to PNG or SVG requires some additional setup in the form of optional dependencies.

I think a clearer model would be to move this functionality to another package, which could then specify (and test) these dependencies itself.


Discussion:
I use Databricks Notebooks and installing browsers and web drivers in a databricks cluster is a lot of overhead to save a chart. Would be great if I could programmatically save them as part of pipeline runs to send out through slack webhooks or imported in dashboards.
Take a look at https://github.com/altair-viz/altair_saver

You can save charts via a nodejs stack instead of a selenium/headless driver stack.
Installing the nodejs stack on a databricks cluster was _not_ easy. But altair_saver worked for me. Thanks!

Resolving patch:
diff --git a/altair/sphinxext/altairgallery.py b/altair/sphinxext/altairgallery.py
--- a/altair/sphinxext/altairgallery.py
+++ b/altair/sphinxext/altairgallery.py
@@ -132,7 +132,7 @@ def save_example_pngs(examples, image_dir, make_thumbnails=True):
                 chart.save(image_file)
                 hashes[filename] = example_hash
             except ImportError:
-                warnings.warn("Could not import selenium: using generic image")
+                warnings.warn("Unable to save image: using generic image")
                 create_generic_image(image_file)
 
             with open(hash_file, 'w') as f:
diff --git a/altair/utils/headless.py b/altair/utils/headless.py
deleted file mode 100644
--- a/altair/utils/headless.py
+++ /dev/null
@@ -1,173 +0,0 @@
-"""
-Utilities that use selenium + chrome headless to save figures
-"""
-
-import contextlib
-import os
-import tempfile
-
-
-@contextlib.contextmanager
-def temporary_filename(**kwargs):
-    """Create and clean-up a temporary file
-
-    Arguments are the same as those passed to tempfile.mkstemp
-
-    We could use tempfile.NamedTemporaryFile here, but that causes issues on
-    windows (see https://bugs.python.org/issue14243).
-    """
-    filedescriptor, filename = tempfile.mkstemp(**kwargs)
-    os.close(filedescriptor)
-
-    try:
-        yield filename
-    finally:
-        if os.path.exists(filename):
-            os.remove(filename)
-
-
-HTML_TEMPLATE = """
-<!DOCTYPE html>
-<html>
-<head>
-  <title>Embedding Vega-Lite</title>
-  <script src="https://cdn.jsdelivr.net/npm/vega@{vega_version}"></script>
-  <script src="https://cdn.jsdelivr.net/npm/vega-lite@{vegalite_version}"></script>
-  <script src="https://cdn.jsdelivr.net/npm/vega-embed@{vegaembed_version}"></script>
-</head>
-<body>
-  <div id="vis"></div>
-</body>
-</html>
-"""
-
-EXTRACT_CODE = {
-'png': """
-        var spec = arguments[0];
-        var mode = arguments[1];
-        var scaleFactor = arguments[2];
-        var done = arguments[3];
-
-        if(mode === 'vega-lite'){
-          // compile vega-lite to vega
-          vegaLite = (typeof vegaLite === "undefined") ? vl : vegaLite;
-          const compiled = vegaLite.compile(spec);
-          spec = compiled.spec;
-        }
-
-        new vega.View(vega.parse(spec), {
-              loader: vega.loader(),
-              logLevel: vega.Warn,
-              renderer: 'none',
-            })
-            .initialize()
-            .toCanvas(scaleFactor)
-            .then(function(canvas){return canvas.toDataURL('image/png');})
-            .then(done)
-            .catch(function(err) { console.error(err); });
-        """,
-'svg': """
-        var spec = arguments[0];
-        var mode = arguments[1];
-        var scaleFactor = arguments[2];
-        var done = arguments[3];
-
-        if(mode === 'vega-lite'){
-          // compile vega-lite to vega
-          vegaLite = (typeof vegaLite === "undefined") ? vl : vegaLite;
-          const compiled = vegaLite.compile(spec);
-          spec = compiled.spec;
-        }
-
-        new vega.View(vega.parse(spec), {
-              loader: vega.loader(),
-              logLevel: vega.Warn,
-              renderer: 'none',
-            })
-            .initialize()
-            .toSVG(scaleFactor)
-            .then(done)
-            .catch(function(err) { console.error(err); });
-        """,
-'vega': """
-        var spec = arguments[0];
-        var mode = arguments[1];
-        var done = arguments[3];
-
-        if(mode === 'vega-lite'){
-          // compile vega-lite to vega
-          vegaLite = (typeof vegaLite === "undefined") ? vl : vegaLite;
-          const compiled = vegaLite.compile(spec);
-          spec = compiled.spec;
-        }
-
-        done(spec);
-        """}
-
-
-def compile_spec(spec, format, mode,
-                 vega_version, vegaembed_version, vegalite_version,
-                 scale_factor=1, driver_timeout=20, webdriver='chrome'):
-    # TODO: detect & use local Jupyter caches of JS packages?
-
-    # selenium is an optional dependency, so import it here
-    try:
-        import selenium.webdriver
-    except ImportError:
-        raise ImportError("selenium package is required "
-                          "for saving chart as {}".format(format))
-
-    if format not in ['png', 'svg', 'vega']:
-        raise NotImplementedError("format must be 'svg', 'png' or 'vega'")
-
-    if mode not in ['vega', 'vega-lite']:
-        raise ValueError("mode must be either 'vega' or 'vega-lite'")
-
-    if vega_version is None:
-        raise ValueError("must specify vega_version")
-
-    if vegaembed_version is None:
-        raise ValueError("must specify vegaembed_version")
-
-    if mode == 'vega-lite' and vegalite_version is None:
-        raise ValueError("must specify vega-lite version")
-
-    if webdriver == 'chrome':
-        webdriver_class = selenium.webdriver.Chrome
-        webdriver_options_class = selenium.webdriver.chrome.options.Options
-    elif webdriver == 'firefox':
-        webdriver_class = selenium.webdriver.Firefox
-        webdriver_options_class = selenium.webdriver.firefox.options.Options
-    else:
-        raise ValueError("webdriver must be 'chrome' or 'firefox'")
-
-    html = HTML_TEMPLATE.format(vega_version=vega_version,
-                                vegalite_version=vegalite_version,
-                                vegaembed_version=vegaembed_version)
-
-    webdriver_options = webdriver_options_class()
-    webdriver_options.add_argument("--headless")
-
-    if issubclass(webdriver_class, selenium.webdriver.Chrome):
-        # for linux/osx root user, need to add --no-sandbox option.
-        # since geteuid doesn't exist on windows, we don't check it
-        if hasattr(os, 'geteuid') and (os.geteuid() == 0):
-            webdriver_options.add_argument('--no-sandbox')
-
-    driver = webdriver_class(options=webdriver_options)
-
-    try:
-        driver.set_page_load_timeout(driver_timeout)
-
-        with temporary_filename(suffix='.html') as htmlfile:
-            with open(htmlfile, 'w') as f:
-                f.write(html)
-            driver.get("file://" + htmlfile)
-            online = driver.execute_script("return navigator.onLine")
-            if not online:
-                raise ValueError("Internet connection required for saving "
-                                 "chart as {}".format(format))
-            return driver.execute_async_script(EXTRACT_CODE[format],
-                                               spec, mode, scale_factor)
-    finally:
-        driver.close()
diff --git a/altair/utils/mimebundle.py b/altair/utils/mimebundle.py
--- a/altair/utils/mimebundle.py
+++ b/altair/utils/mimebundle.py
@@ -1,6 +1,3 @@
-import base64
-
-from .headless import compile_spec
 from .html import spec_to_html
 
 
@@ -12,13 +9,13 @@ def spec_to_mimebundle(spec, format, mode=None,
     """Convert a vega/vega-lite specification to a mimebundle
 
     The mimebundle type is controlled by the ``format`` argument, which can be
-    one of the following ['png', 'svg', 'vega', 'vega-lite', 'html', 'json']
+    one of the following ['html', 'json', 'png', 'svg', 'pdf', 'vega', 'vega-lite']
 
     Parameters
     ----------
     spec : dict
         a dictionary representing a vega-lite plot spec
-    format : string {'png', 'svg', 'vega', 'vega-lite', 'html', 'json'}
+    format : string {'html', 'json', 'png', 'svg', 'pdf', 'vega', 'vega-lite'}
         the file format to be saved.
     mode : string {'vega', 'vega-lite'}
         The rendering mode.
@@ -38,9 +35,8 @@ def spec_to_mimebundle(spec, format, mode=None,
 
     Note
     ----
-    The png, svg, and vega outputs require the pillow and selenium Python modules
-    to be installed. Additionally they requires either chromedriver
-    (if webdriver=='chrome') or geckodriver (if webdriver=='firefox')
+    The png, svg, pdf, and vega outputs require the altair_saver package
+    to be installed.
     """
     if mode not in ['vega', 'vega-lite']:
         raise ValueError("mode must be either 'vega' or 'vega-lite'")
@@ -49,34 +45,29 @@ def spec_to_mimebundle(spec, format, mode=None,
         if vega_version is None:
             raise ValueError("Must specify vega_version")
         return {'application/vnd.vega.v{}+json'.format(vega_version[0]): spec}
-    elif format in ['png', 'svg', 'vega']:
-        render = compile_spec(spec, format=format, mode=mode,
-                              vega_version=vega_version,
-                              vegaembed_version=vegaembed_version,
-                              vegalite_version=vegalite_version, **kwargs)
-        if format == 'png':
-            render = base64.b64decode(render.split(',', 1)[1].encode())
-            return {'image/png': render}
-        elif format == 'svg':
-            return {'image/svg+xml': render}
-        elif format == 'vega':
-            assert mode == 'vega-lite'  # TODO: handle vega->vega conversion more gracefully
-            return {'application/vnd.vega.v{}+json'.format(vega_version[0]): render}
-    elif format == 'html':
+    if format in ['png', 'svg', 'pdf', 'vega']:
+        try:
+            import altair_saver
+        except ImportError:
+            raise ValueError(
+                "Saving charts in {fmt!r} format requires the altair_saver package: "
+                "see http://github.com/altair-viz/altair_saver/".format(fmt=format)
+            )
+        return altair_saver.render(spec, format, mode=mode, **kwargs)
+    if format == 'html':
         html = spec_to_html(spec, mode=mode,
                             vega_version=vega_version,
                             vegaembed_version=vegaembed_version,
                             vegalite_version=vegalite_version, **kwargs)
         return {'text/html': html}
-    elif format == 'vega-lite':
+    if format == 'vega-lite':
         assert mode == 'vega-lite'  # sanity check: should never be False
         if mode == 'vega':
             raise ValueError("Cannot convert a vega spec to vegalite")
         if vegalite_version is None:
             raise ValueError("Must specify vegalite_version")
         return {'application/vnd.vegalite.v{}+json'.format(vegalite_version[0]): spec}
-    elif format == 'json':
+    if format == 'json':
         return {'application/json': spec}
-    else:
-        raise ValueError("format must be one of "
-                         "['png', 'svg', 'vega', 'vega-lite', 'html', 'json']")
+    raise ValueError("format must be one of "
+                     "['html', 'json', 'png', 'svg', 'pdf', 'vega', 'vega-lite']")


Test patch:
diff --git a/altair/utils/tests/test_mimebundle.py b/altair/utils/tests/test_mimebundle.py
--- a/altair/utils/tests/test_mimebundle.py
+++ b/altair/utils/tests/test_mimebundle.py
@@ -1,195 +1,192 @@
 import pytest
-import json
-
-try:
-    import selenium
-except ImportError:
-    selenium = None
 
+import altair as alt
 from ..mimebundle import spec_to_mimebundle
 
 
-# example from https://vega.github.io/editor/#/examples/vega-lite/bar
-VEGALITE_SPEC = json.loads(
-    """
-        {
-    "$schema": "https://vega.github.io/schema/vega-lite/v2.json",
-    "description": "A simple bar chart with embedded data.",
-    "data": {
-        "values": [
-        {"a": "A","b": 28}, {"a": "B","b": 55}, {"a": "C","b": 43},
-        {"a": "D","b": 91}, {"a": "E","b": 81}, {"a": "F","b": 53},
-        {"a": "G","b": 19}, {"a": "H","b": 87}, {"a": "I","b": 52}
-        ]
-    },
-    "mark": "bar",
-    "encoding": {
-        "x": {"field": "a", "type": "ordinal"},
-        "y": {"field": "b", "type": "quantitative"}
-    }
-    }
-    """
-)
+def require_altair_saver(func):
+    try:
+        import altair_saver  # noqa: F401
+    except ImportError:
+        return pytest.mark.skip("altair_saver not importable; cannot run saver tests")(func)
+    else:
+        return func
 
-VEGA_SPEC = json.loads(
-    """
-    {
-    "$schema": "https://vega.github.io/schema/vega/v3.0.json",
-    "description": "A simple bar chart with embedded data.",
-    "autosize": "pad",
-    "padding": 5,
-    "height": 200,
-    "style": "cell",
-    "data": [
-        {
-        "name": "source_0",
-        "values": [
-            {"a": "A", "b": 28},
-            {"a": "B", "b": 55},
-            {"a": "C", "b": 43},
-            {"a": "D", "b": 91},
-            {"a": "E", "b": 81},
-            {"a": "F", "b": 53},
-            {"a": "G", "b": 19},
-            {"a": "H", "b": 87},
-            {"a": "I", "b": 52}
-        ]
+
+@pytest.fixture
+def vegalite_spec():
+    return {
+        "$schema": "https://vega.github.io/schema/vega-lite/v4.json",
+        "description": "A simple bar chart with embedded data.",
+        "data": {
+            "values": [
+                {"a": "A", "b": 28},
+                {"a": "B", "b": 55},
+                {"a": "C", "b": 43},
+                {"a": "D", "b": 91},
+                {"a": "E", "b": 81},
+                {"a": "F", "b": 53},
+                {"a": "G", "b": 19},
+                {"a": "H", "b": 87},
+                {"a": "I", "b": 52},
+            ]
         },
-        {
-        "name": "data_0",
-        "source": "source_0",
-        "transform": [
-            {"type": "formula", "expr": "toNumber(datum[\\"b\\"])", "as": "b"},
-            {
-            "type": "filter",
-            "expr": "datum[\\"b\\"] !== null && !isNaN(datum[\\"b\\"])"
-            }
-        ]
-        }
-    ],
-    "signals": [
-        {"name": "x_step", "value": 21},
-        {
-        "name": "width",
-        "update": "bandspace(domain('x').length, 0.1, 0.05) * x_step"
-        }
-    ],
-    "marks": [
-        {
-        "name": "marks",
-        "type": "rect",
-        "style": ["bar"],
-        "from": {"data": "data_0"},
-        "encode": {
-            "update": {
-            "fill": {"value": "#4c78a8"},
-            "x": {"scale": "x", "field": "a"},
-            "width": {"scale": "x", "band": true},
-            "y": {"scale": "y", "field": "b"},
-            "y2": {"scale": "y", "value": 0}
-            }
-        }
-        }
-    ],
-    "scales": [
-        {
-        "name": "x",
-        "type": "band",
-        "domain": {"data": "data_0", "field": "a", "sort": true},
-        "range": {"step": {"signal": "x_step"}},
-        "paddingInner": 0.1,
-        "paddingOuter": 0.05
+        "mark": "bar",
+        "encoding": {
+            "x": {"field": "a", "type": "ordinal"},
+            "y": {"field": "b", "type": "quantitative"},
         },
-        {
-        "name": "y",
-        "type": "linear",
-        "domain": {"data": "data_0", "field": "b"},
-        "range": [{"signal": "height"}, 0],
-        "nice": true,
-        "zero": true
-        }
-    ],
-    "axes": [
-        {
-        "scale": "x",
-        "orient": "bottom",
-        "title": "a",
-        "labelOverlap": true,
-        "encode": {
-            "labels": {
-            "update": {
-                "angle": {"value": 270},
-                "align": {"value": "right"},
-                "baseline": {"value": "middle"}
-            }
+    }
+
+
+@pytest.fixture
+def vega_spec():
+    return {
+        "$schema": "https://vega.github.io/schema/vega/v5.json",
+        "axes": [
+            {
+                "domain": False,
+                "grid": True,
+                "gridScale": "x",
+                "labels": False,
+                "maxExtent": 0,
+                "minExtent": 0,
+                "orient": "left",
+                "scale": "y",
+                "tickCount": {"signal": "ceil(height/40)"},
+                "ticks": False,
+                "zindex": 0,
+            },
+            {
+                "grid": False,
+                "labelAlign": "right",
+                "labelAngle": 270,
+                "labelBaseline": "middle",
+                "labelOverlap": True,
+                "orient": "bottom",
+                "scale": "x",
+                "title": "a",
+                "zindex": 0,
+            },
+            {
+                "grid": False,
+                "labelOverlap": True,
+                "orient": "left",
+                "scale": "y",
+                "tickCount": {"signal": "ceil(height/40)"},
+                "title": "b",
+                "zindex": 0,
+            },
+        ],
+        "background": "white",
+        "data": [
+            {
+                "name": "source_0",
+                "values": [
+                    {"a": "A", "b": 28},
+                    {"a": "B", "b": 55},
+                    {"a": "C", "b": 43},
+                    {"a": "D", "b": 91},
+                    {"a": "E", "b": 81},
+                    {"a": "F", "b": 53},
+                    {"a": "G", "b": 19},
+                    {"a": "H", "b": 87},
+                    {"a": "I", "b": 52},
+                ],
+            },
+            {
+                "name": "data_0",
+                "source": "source_0",
+                "transform": [
+                    {
+                        "expr": 'isValid(datum["b"]) && isFinite(+datum["b"])',
+                        "type": "filter",
+                    }
+                ],
+            },
+        ],
+        "description": "A simple bar chart with embedded data.",
+        "height": 200,
+        "marks": [
+            {
+                "encode": {
+                    "update": {
+                        "fill": {"value": "#4c78a8"},
+                        "width": {"band": True, "scale": "x"},
+                        "x": {"field": "a", "scale": "x"},
+                        "y": {"field": "b", "scale": "y"},
+                        "y2": {"scale": "y", "value": 0},
+                    }
+                },
+                "from": {"data": "data_0"},
+                "name": "marks",
+                "style": ["bar"],
+                "type": "rect",
             }
-        },
-        "zindex": 1
-        },
-        {
-        "scale": "y",
-        "orient": "left",
-        "title": "b",
-        "labelOverlap": true,
-        "tickCount": {"signal": "ceil(height/40)"},
-        "zindex": 1
-        },
-        {
-        "scale": "y",
-        "orient": "left",
-        "grid": true,
-        "tickCount": {"signal": "ceil(height/40)"},
-        "gridScale": "x",
-        "domain": false,
-        "labels": false,
-        "maxExtent": 0,
-        "minExtent": 0,
-        "ticks": false,
-        "zindex": 0
-        }
-    ],
-    "config": {"axisY": {"minExtent": 30}}
+        ],
+        "padding": 5,
+        "scales": [
+            {
+                "domain": {"data": "data_0", "field": "a", "sort": True},
+                "name": "x",
+                "paddingInner": 0.1,
+                "paddingOuter": 0.05,
+                "range": {"step": {"signal": "x_step"}},
+                "type": "band",
+            },
+            {
+                "domain": {"data": "data_0", "field": "b"},
+                "name": "y",
+                "nice": True,
+                "range": [{"signal": "height"}, 0],
+                "type": "linear",
+                "zero": True,
+            },
+        ],
+        "signals": [
+            {"name": "x_step", "value": 20},
+            {
+                "name": "width",
+                "update": "bandspace(domain('x').length, 0.1, 0.05) * x_step",
+            },
+        ],
+        "style": "cell",
     }
-    """
-)
 
-VEGAEMBED_VERSION = '3.14.0'
-VEGALITE_VERSION = '2.3.1'
-VEGA_VERSION = '3.3.1'
 
-@pytest.mark.skipif('not selenium')
-def test_spec_to_vega_mimebundle():
-    try:
-        bundle = spec_to_mimebundle(
-            spec=VEGALITE_SPEC,
-            format='vega',
-            mode='vega-lite',
-            vega_version=VEGA_VERSION,
-            vegalite_version=VEGALITE_VERSION,
-            vegaembed_version=VEGAEMBED_VERSION
-        )
-    except ValueError as err:
-        if str(err).startswith('Internet connection'):
-            pytest.skip("web connection required for png/svg export")
-        else:
-            raise
-    assert bundle == {'application/vnd.vega.v3+json': VEGA_SPEC}
+@require_altair_saver
+def test_vegalite_to_vega_mimebundle(vegalite_spec, vega_spec):
+    bundle = spec_to_mimebundle(
+        spec=vegalite_spec,
+        format="vega",
+        mode="vega-lite",
+        vega_version=alt.VEGA_VERSION,
+        vegalite_version=alt.VEGALITE_VERSION,
+        vegaembed_version=alt.VEGAEMBED_VERSION,
+    )
+    assert bundle == {"application/vnd.vega.v5+json": vega_spec}
 
 
-def test_spec_to_vegalite_mimebundle():
+def test_spec_to_vegalite_mimebundle(vegalite_spec):
     bundle = spec_to_mimebundle(
-        spec=VEGALITE_SPEC,
-        mode='vega-lite',
-        format='vega-lite',
-        vegalite_version=VEGALITE_VERSION
+        spec=vegalite_spec,
+        mode="vega-lite",
+        format="vega-lite",
+        vegalite_version=alt.VEGALITE_VERSION,
     )
-    assert bundle == {'application/vnd.vegalite.v2+json': VEGALITE_SPEC}
+    assert bundle == {"application/vnd.vegalite.v4+json": vegalite_spec}
 
 
-def test_spec_to_json_mimebundle():
+def test_spec_to_vega_mimebundle(vega_spec):
     bundle = spec_to_mimebundle(
-        spec=VEGALITE_SPEC,
-        mode='vega-lite',
-        format='json',
+        spec=vega_spec,
+        mode="vega",
+        format="vega",
+        vega_version=alt.VEGA_VERSION
     )
-    assert bundle == {'application/json': VEGALITE_SPEC}
+    assert bundle == {"application/vnd.vega.v5+json": vega_spec}
+
+
+def test_spec_to_json_mimebundle():
+    bundle = spec_to_mimebundle(spec=vegalite_spec, mode="vega-lite", format="json",)
+    assert bundle == {"application/json": vegalite_spec}
diff --git a/altair/vegalite/v3/tests/test_api.py b/altair/vegalite/v3/tests/test_api.py
--- a/altair/vegalite/v3/tests/test_api.py
+++ b/altair/vegalite/v3/tests/test_api.py
@@ -14,9 +14,9 @@
 from altair.utils import AltairDeprecationWarning
 
 try:
-    import selenium
+    import altair_saver  # noqa: F401
 except ImportError:
-    selenium = None
+    altair_saver = None
 
 
 def getargs(*args, **kwargs):
@@ -230,44 +230,39 @@ def test_selection_expression():
 
 
 @pytest.mark.parametrize('format', ['html', 'json', 'png', 'svg'])
-@pytest.mark.skipif('not selenium')
 def test_save(format, basic_chart):
-    if format in ['html', 'json', 'svg']:
-        out = io.StringIO()
-        mode = 'r'
-    else:
+    if format == 'png':
         out = io.BytesIO()
         mode = 'rb'
+    else:
+        out = io.StringIO()
+        mode = 'r'
+
+    if format in ['svg', 'png'] and not altair_saver:
+        with pytest.raises(ValueError) as err:
+            basic_chart.save(out, format=format)
+        assert "github.com/altair-viz/altair_saver" in str(err.value)
+        return
+
+    basic_chart.save(out, format=format)
+    out.seek(0)
+    content = out.read()
+
+    if format == 'json':
+        assert '$schema' in json.loads(content)
+    if format == 'html':
+        assert content.startswith('<!DOCTYPE html>')
     
     fid, filename = tempfile.mkstemp(suffix='.' + format)
     os.close(fid)
-    
+
     try:
-        try:
-            basic_chart.save(out, format=format)
-            basic_chart.save(filename)
-        except ValueError as err:
-            if str(err).startswith('Internet connection'):
-                pytest.skip("web connection required for png/svg export")
-            else:
-                raise
-
-        out.seek(0)
+        basic_chart.save(filename)
         with open(filename, mode) as f:
-            assert f.read() == out.read()
+            assert f.read() == content
     finally:
         os.remove(filename)
 
-    out.seek(0)
-
-    if format == 'json':
-        spec = json.load(out)
-        assert '$schema' in spec
-
-    elif format == 'html':
-        content = out.read()
-        assert content.startswith('<!DOCTYPE html>')
-
 
 def test_facet():
     # wrapped facet
diff --git a/altair/vegalite/v4/tests/test_api.py b/altair/vegalite/v4/tests/test_api.py
--- a/altair/vegalite/v4/tests/test_api.py
+++ b/altair/vegalite/v4/tests/test_api.py
@@ -13,9 +13,9 @@
 import altair.vegalite.v4 as alt
 
 try:
-    import selenium
+    import altair_saver  # noqa: F401
 except ImportError:
-    selenium = None
+    altair_saver = None
 
 
 def getargs(*args, **kwargs):
@@ -229,44 +229,39 @@ def test_selection_expression():
 
 
 @pytest.mark.parametrize('format', ['html', 'json', 'png', 'svg'])
-@pytest.mark.skipif('not selenium')
 def test_save(format, basic_chart):
-    if format in ['html', 'json', 'svg']:
-        out = io.StringIO()
-        mode = 'r'
-    else:
+    if format == 'png':
         out = io.BytesIO()
         mode = 'rb'
+    else:
+        out = io.StringIO()
+        mode = 'r'
+
+    if format in ['svg', 'png'] and not altair_saver:
+        with pytest.raises(ValueError) as err:
+            basic_chart.save(out, format=format)
+        assert "github.com/altair-viz/altair_saver" in str(err.value)
+        return
+
+    basic_chart.save(out, format=format)
+    out.seek(0)
+    content = out.read()
+
+    if format == 'json':
+        assert '$schema' in json.loads(content)
+    if format == 'html':
+        assert content.startswith('<!DOCTYPE html>')
     
     fid, filename = tempfile.mkstemp(suffix='.' + format)
     os.close(fid)
 
     try:
-        try:
-            basic_chart.save(out, format=format)
-            basic_chart.save(filename)
-        except ValueError as err:
-            if str(err).startswith('Internet connection'):
-                pytest.skip("web connection required for png/svg export")
-            else:
-                raise
-
-        out.seek(0)
+        basic_chart.save(filename)
         with open(filename, mode) as f:
-            assert f.read() == out.read()
+            assert f.read() == content
     finally:
         os.remove(filename)
 
-    out.seek(0)
-
-    if format == 'json':
-        spec = json.load(out)
-        assert '$schema' in spec
-
-    elif format == 'html':
-        content = out.read()
-        assert content.startswith('<!DOCTYPE html>')
-
 
 def test_facet():
     # wrapped facet