Test vl-convert engine in chart.save test (#2784)

* Explicitly pass engine argument through save function

* Update test_save to test both altair_saver and vl-convert engiens

* flake8

altair/utils/mimebundle.py

@@ -9,6 +9,7 @@ def spec_to_mimebundle(
     vega_version=None,
     vegaembed_version=None,
     vegalite_version=None,
+    engine=None,
     **kwargs,
 ):
     """Convert a vega/vega-lite specification to a mimebundle
@@ -30,6 +31,8 @@ def spec_to_mimebundle(
         The version of vegaembed.js to use
     vegalite_version : string
         The version of vegalite.js to use. Only required if mode=='vega-lite'
+    engine: string {'vl-convert', 'altair_saver'}
+        the conversion engine to use for 'png', 'svg', 'pdf', and 'vega' formats
     **kwargs :
         Additional arguments will be passed to the generating function
 
@@ -51,7 +54,9 @@ def spec_to_mimebundle(
             raise ValueError("Must specify vega_version")
         return {"application/vnd.vega.v{}+json".format(vega_version[0]): spec}
     if format in ["png", "svg", "pdf", "vega"]:
-        return _spec_to_mimebundle_with_engine(spec, format, mode, **kwargs)
+        return _spec_to_mimebundle_with_engine(
+            spec, format, mode, engine=engine, **kwargs
+        )
     if format == "html":
         html = spec_to_html(
             spec,

altair/utils/save.py

@@ -26,6 +26,7 @@ def save(
     json_kwds=None,
     webdriver=None,
     scale_factor=1,
+    engine=None,
     **kwargs,
 ):
     """Save a chart to file in a variety of formats
@@ -61,6 +62,8 @@ def save(
         Webdriver to use for png or svg output
     scale_factor : float (optional)
         scale_factor to use to change size/resolution of png or svg output
+    engine: string {'vl-convert', 'altair_saver'}
+        the conversion engine to use for 'png', 'svg', and 'pdf' formats
     **kwargs :
         additional kwargs passed to spec_to_mimebundle.
     """
@@ -122,6 +125,7 @@ def save(
             vegaembed_version=vegaembed_version,
             webdriver=webdriver,
             scale_factor=scale_factor,
+            engine=engine,
             **kwargs,
         )
         if format == "png":
@@ -131,4 +135,4 @@ def save(
         else:
             write_file_or_filename(fp, mimebundle["image/svg+xml"], mode="w")
     else:
-        raise ValueError("unrecognized format: '{}'".format(format))
+        raise ValueError("Unsupported format: '{}'".format(format))

altair/vegalite/v5/tests/test_api.py

@@ -18,6 +18,11 @@
 except ImportError:
     altair_saver = None
 
+try:
+    import vl_convert as vlc  # noqa: F401
+except ImportError:
+    vlc = None
+
 
 def getargs(*args, **kwargs):
     return args, kwargs
@@ -253,45 +258,76 @@ def test_selection_expression():
         selection.__magic__
 
 
-@pytest.mark.parametrize("format", ["html", "json", "png", "svg", "pdf"])
-def test_save(format, basic_chart):
+@pytest.mark.parametrize("format", ["html", "json", "png", "svg", "pdf", "bogus"])
+@pytest.mark.parametrize("engine", ["altair_saver", "vl-convert"])
+def test_save(format, engine, basic_chart):
     if format in ["pdf", "png"]:
         out = io.BytesIO()
         mode = "rb"
     else:
         out = io.StringIO()
         mode = "r"
 
-    if format in ["svg", "png", "pdf"]:
-        if 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
-        elif format not in altair_saver.available_formats():
-            with pytest.raises(ValueError) as err:
-                basic_chart.save(out, format=format)
-            assert f"No enabled saver found that supports format='{format}'" in str(
-                err.value
-            )
-            return
-
-    basic_chart.save(out, format=format)
+    if format in ["svg", "png", "pdf", "bogus"]:
+        if engine == "altair_saver":
+            if altair_saver is None:
+                with pytest.raises(ValueError) as err:
+                    basic_chart.save(out, format=format, engine=engine)
+                assert "altair_saver" in str(err.value)
+                return
+            elif format == "bogus":
+                with pytest.raises(ValueError) as err:
+                    basic_chart.save(out, format=format, engine=engine)
+                assert f"Unsupported format: '{format}'" in str(err.value)
+                return
+            elif format not in altair_saver.available_formats():
+                with pytest.raises(ValueError) as err:
+                    basic_chart.save(out, format=format, engine=engine)
+                assert f"No enabled saver found that supports format='{format}'" in str(
+                    err.value
+                )
+                return
+
+        elif engine == "vl-convert":
+            if vlc is None:
+                with pytest.raises(ValueError) as err:
+                    basic_chart.save(out, format=format, engine=engine)
+                assert "vl-convert-python" in str(err.value)
+                return
+            elif format == "pdf":
+                with pytest.raises(ValueError) as err:
+                    basic_chart.save(out, format=format, engine=engine)
+                assert (
+                    f"The 'vl-convert' conversion engine does not support the '{format}' format"
+                    in str(err.value)
+                )
+                return
+            elif format not in ("png", "svg"):
+                with pytest.raises(ValueError) as err:
+                    basic_chart.save(out, format=format, engine=engine)
+                assert f"Unsupported format: '{format}'" in str(err.value)
+                return
+
+    basic_chart.save(out, format=format, engine=engine)
     out.seek(0)
     content = out.read()
 
     if format == "json":
         assert "$schema" in json.loads(content)
-    if format == "html":
+    elif format == "html":
         assert content.startswith("<!DOCTYPE html>")
+    elif format == "svg":
+        assert content.startswith("<svg")
+    elif format == "png":
+        assert content.startswith(b"\x89PNG")
 
     fid, filename = tempfile.mkstemp(suffix="." + format)
     os.close(fid)
 
     # test both string filenames and pathlib.Paths
     for fp in [filename, pathlib.Path(filename)]:
         try:
-            basic_chart.save(fp)
+            basic_chart.save(fp, format=format, engine=engine)
             with open(fp, mode) as f:
                 assert f.read()[:1000] == content[:1000]
         finally: