fix: Repair vega export (#3889)

* Repair vega export

Even if not advertise, it is better to not fail for vega by going in the svg case

* improve fix and add test

* ruff changes

* fix ruff, reduce complexity

* fix mypy

---------

Co-authored-by: Mattijn van Hoek <[email protected]>

Author: bobot

altair/utils/save.py

@@ -75,12 +75,105 @@ def set_inspect_mode_argument(
     return mode
 
 
+def _save_mimebundle_format(
+    spec: dict[str, Any],
+    format: Literal["html", "png", "svg", "pdf", "vega"],
+    fp: str | Path | IO,
+    inner_mode: Literal["vega-lite"],
+    vega_version: str | None,
+    vegalite_version: str | None,
+    vegaembed_version: str | None,
+    embed_options: dict[str, Any] | None,
+    json_kwds: dict[str, Any],
+    encoding: str,
+    scale_factor: float,
+    engine: Literal["vl-convert"] | None,
+    inline: bool,
+    **kwargs: Any,
+) -> None:
+    """Save chart using spec_to_mimebundle for formats that require it."""
+    if format == "html":
+        if inline:
+            kwargs["template"] = "inline"
+        mb_result: dict[str, str] = spec_to_mimebundle(
+            spec=spec,
+            format=format,
+            mode=inner_mode,
+            vega_version=vega_version,
+            vegalite_version=vegalite_version,
+            vegaembed_version=vegaembed_version,
+            embed_options=embed_options,
+            json_kwds=json_kwds,
+            **kwargs,
+        )
+        write_file_or_filename(fp, mb_result["text/html"], mode="w", encoding=encoding)
+    elif format == "png":
+        mb_result_png: tuple[dict[str, Any], dict[str, Any]] = spec_to_mimebundle(
+            spec=spec,
+            format=format,
+            mode=inner_mode,
+            vega_version=vega_version,
+            vegalite_version=vegalite_version,
+            vegaembed_version=vegaembed_version,
+            embed_options=embed_options,
+            scale_factor=scale_factor,
+            engine=engine,
+            **kwargs,
+        )
+        write_file_or_filename(fp, mb_result_png[0]["image/png"], mode="wb")
+    elif format == "svg":
+        mb_result = spec_to_mimebundle(
+            spec=spec,
+            format=format,
+            mode=inner_mode,
+            vega_version=vega_version,
+            vegalite_version=vegalite_version,
+            vegaembed_version=vegaembed_version,
+            embed_options=embed_options,
+            scale_factor=scale_factor,
+            engine=engine,
+            **kwargs,
+        )
+        write_file_or_filename(
+            fp, mb_result["image/svg+xml"], mode="w", encoding=encoding
+        )
+    elif format == "pdf":
+        mb_result = spec_to_mimebundle(
+            spec=spec,
+            format=format,
+            mode=inner_mode,
+            vega_version=vega_version,
+            vegalite_version=vegalite_version,
+            vegaembed_version=vegaembed_version,
+            embed_options=embed_options,
+            scale_factor=scale_factor,
+            engine=engine,
+            **kwargs,
+        )
+        write_file_or_filename(fp, mb_result["application/pdf"], mode="wb")
+    else:  # vega
+        mb_result = spec_to_mimebundle(
+            spec=spec,
+            format=format,
+            mode=inner_mode,
+            vega_version=vega_version,
+            vegalite_version=vegalite_version,
+            vegaembed_version=vegaembed_version,
+            embed_options=embed_options,
+            scale_factor=scale_factor,
+            engine=engine,
+            **kwargs,
+        )
+        json_spec = json.dumps(mb_result["application/vnd.vega.v6+json"], **json_kwds)
+        write_file_or_filename(fp, json_spec, mode="w", encoding=encoding)
+
+
 def save(
     chart,
     fp: str | Path | IO,
     vega_version: str | None,
     vegaembed_version: str | None,
-    format: Literal["json", "html", "png", "svg", "pdf"] | None = None,
+    format: Literal["json", "html", "png", "svg", "pdf", "vega"] | None = None,
     mode: Literal["vega-lite"] | None = None,
     vegalite_version: str | None = None,
     embed_options: dict | None = None,
@@ -93,7 +186,7 @@ def save(
     """
     Save a chart to file in a variety of formats.
 
-    Supported formats are [json, html, png, svg, pdf]
+    Supported formats are [json, html, png, svg, pdf, vega]
 
     Parameters
     ----------
@@ -102,7 +195,7 @@ def save(
     fp : string filename, pathlib.Path or file-like object
         file to which to write the chart.
     format : string (optional)
-        the format to write: one of ['json', 'html', 'png', 'svg', 'pdf'].
+        the format to write: one of ['json', 'html', 'png', 'svg', 'pdf', 'vega'].
         If not specified, the format will be determined from the filename.
     mode : string (optional)
         Must be 'vega-lite'. If not specified, then infer the mode from
@@ -123,7 +216,7 @@ def save(
     scale_factor : float (optional)
         scale_factor to use to change size/resolution of png or svg output
     engine: string {'vl-convert'}
-        the conversion engine to use for 'png', 'svg', and 'pdf' formats
+        the conversion engine to use for 'png', 'svg', 'pdf', and 'vega' formats
     inline: bool (optional)
         If False (default), the required JavaScript libraries are loaded
         from a CDN location in the resulting html file.
@@ -154,56 +247,23 @@ def perform_save() -> None:
         if format == "json":
             json_spec = json.dumps(spec, **json_kwds)
             write_file_or_filename(fp, json_spec, mode="w", encoding=encoding)
-        elif format == "html":
-            if inline:
-                kwargs["template"] = "inline"
-            mb_html = spec_to_mimebundle(
+        elif format in {"html", "png", "svg", "pdf", "vega"}:
+            _save_mimebundle_format(
                 spec=spec,
                 format=format,
-                mode=inner_mode,
+                fp=fp,
+                inner_mode=inner_mode,
                 vega_version=vega_version,
                 vegalite_version=vegalite_version,
                 vegaembed_version=vegaembed_version,
                 embed_options=embed_options,
                 json_kwds=json_kwds,
-                **kwargs,
-            )
-            write_file_or_filename(
-                fp, mb_html["text/html"], mode="w", encoding=encoding
-            )
-        elif format == "png":
-            mb_png = spec_to_mimebundle(
-                spec=spec,
-                format=format,
-                mode=inner_mode,
-                vega_version=vega_version,
-                vegalite_version=vegalite_version,
-                vegaembed_version=vegaembed_version,
-                embed_options=embed_options,
-                scale_factor=scale_factor,
-                engine=engine,
-                **kwargs,
-            )
-            write_file_or_filename(fp, mb_png[0]["image/png"], mode="wb")
-        elif format in {"svg", "pdf", "vega"}:
-            mb_any = spec_to_mimebundle(
-                spec=spec,
-                format=format,
-                mode=inner_mode,
-                vega_version=vega_version,
-                vegalite_version=vegalite_version,
-                vegaembed_version=vegaembed_version,
-                embed_options=embed_options,
+                encoding=encoding,
                 scale_factor=scale_factor,
                 engine=engine,
+                inline=inline,
                 **kwargs,
             )
-            if format == "pdf":
-                write_file_or_filename(fp, mb_any["application/pdf"], mode="wb")
-            else:
-                write_file_or_filename(
-                    fp, mb_any["image/svg+xml"], mode="w", encoding=encoding
-                )
         else:
             msg = f"Unsupported format: '{format}'"
             raise ValueError(msg)

tests/vegalite/v6/test_api.py

@@ -825,7 +825,9 @@ def test_selection_expression():
         getattr(selection, magic_attr)
 
 
-@pytest.mark.parametrize("format", ["html", "json", "png", "svg", "pdf", "bogus"])
+@pytest.mark.parametrize(
+    "format", ["html", "json", "png", "svg", "pdf", "vega", "bogus"]
+)
 @pytest.mark.parametrize("engine", ["vl-convert"])
 def test_save(format, engine, basic_chart):
     if format in {"pdf", "png"}:
@@ -835,7 +837,7 @@ def test_save(format, engine, basic_chart):
         out = io.StringIO()
         mode = "r"
 
-    if format in {"svg", "png", "pdf", "bogus"} and engine == "vl-convert":
+    if format in {"svg", "png", "pdf", "vega", "bogus"} and engine == "vl-convert":
         if format == "bogus":
             with pytest.raises(ValueError) as err:  # noqa: PT011
                 basic_chart.save(out, format=format, engine=engine)
@@ -853,6 +855,15 @@ def test_save(format, engine, basic_chart):
 
     if format == "json":
         assert "$schema" in json.loads(content)
+    elif format == "vega":
+        vega_spec = json.loads(content)
+        assert "$schema" in vega_spec
+        # Verify it's a Vega schema, not Vega-Lite
+        assert "vega/v" in vega_spec["$schema"]
+        assert "vega-lite" not in vega_spec["$schema"]
+        # Verify it has Vega-specific structures
+        assert "marks" in vega_spec
+        assert "scales" in vega_spec
     elif format == "html":
         assert isinstance(content, str)
         assert content.startswith("<!DOCTYPE html>")