Repository: altair-viz/altair
Pull request #3071 (2023-05-27 14:28:29)

Add Vega-Lite to Vega compiler registry and format arg to to_dict() and to_json()
This PR is a first step towards the proposal discussed in https://github.com/altair-viz/altair/discussions/3054 to add a `chart.transformed_data()` function that relies on VegaFusion.

Because VegaFusion works on Vega specs, it's necessary to compile the Vega-Lite spec produced by `chart.to_dict()` into a Vega spec.  This PR adds a new `chart.to_vega()` method to perform this compilation.  Vega-Lite to Vega compilation support is provided by vl-convert.

To maximize future flexibility, this PR includes a new `PluginRegistry` for functions that compile Vega-Lite to Vega and registers a default implementation using vl-convert. My thinking is that this approach will make it possible to support the `to_vega()` method in contexts where vl-convert isn't available. For example, when running in pyodide, a custom plugin could be registered that uses the Vega-Lite JavaScript library directly. Or if running on an architecture that `vl-convert-python` isn't compiled for, a custom implementation could be registered that makes an API call to a remote server running vl-convert.

A couple of tests are added in tests/utils/test_compiler.py


Patch:
diff --git a/altair/__init__.py b/altair/__init__.py
index ca4ee1af..1a66239a 100644
--- a/altair/__init__.py
+++ b/altair/__init__.py
@@ -556,6 +556,7 @@ __all__ = [
     "binding_select",
     "channels",
     "check_fields_and_encodings",
+    "compiler",
     "concat",
     "condition",
     "core",
@@ -596,6 +597,7 @@ __all__ = [
     "value",
     "vconcat",
     "vegalite",
+    "vegalite_compilers",
     "with_property_setters",
 ]
 
diff --git a/altair/utils/compiler.py b/altair/utils/compiler.py
new file mode 100644
index 00000000..0944c92f
--- /dev/null
+++ b/altair/utils/compiler.py
@@ -0,0 +1,11 @@
+from typing import Callable
+from altair.utils import PluginRegistry
+
+# ==============================================================================
+# Vega-Lite to Vega compiler registry
+# ==============================================================================
+VegaLiteCompilerType = Callable[[dict], dict]
+
+
+class VegaLiteCompilerRegistry(PluginRegistry[VegaLiteCompilerType]):
+    pass
diff --git a/altair/utils/schemapi.py b/altair/utils/schemapi.py
index 4bb04cb2..84351b9b 100644
--- a/altair/utils/schemapi.py
+++ b/altair/utils/schemapi.py
@@ -737,6 +737,7 @@ class SchemaBase:
     def to_dict(
         self,
         validate: bool = True,
+        *,
         ignore: Optional[List[str]] = None,
         context: Optional[Dict[str, Any]] = None,
     ) -> dict:
@@ -829,6 +830,7 @@ class SchemaBase:
         validate: bool = True,
         indent: int = 2,
         sort_keys: bool = True,
+        *,
         ignore: Optional[List[str]] = None,
         context: Optional[Dict[str, Any]] = None,
         **kwargs,
diff --git a/altair/vegalite/v5/__init__.py b/altair/vegalite/v5/__init__.py
index 38202d89..8c75d5cc 100644
--- a/altair/vegalite/v5/__init__.py
+++ b/altair/vegalite/v5/__init__.py
@@ -5,6 +5,7 @@ from .api import *
 from ...expr import datum, expr  # type: ignore[no-redef]
 
 from .display import VegaLite, renderers
+from .compiler import vegalite_compilers
 
 from .data import (
     MaxRowsError,
diff --git a/altair/vegalite/v5/api.py b/altair/vegalite/v5/api.py
index 8418b79e..33844a6a 100644
--- a/altair/vegalite/v5/api.py
+++ b/altair/vegalite/v5/api.py
@@ -20,6 +20,7 @@ from .data import data_transformers
 from ... import utils, expr
 from .display import renderers, VEGALITE_VERSION, VEGAEMBED_VERSION, VEGA_VERSION
 from .theme import themes
+from .compiler import vegalite_compilers
 
 if sys.version_info >= (3, 11):
     from typing import Self
@@ -819,6 +820,8 @@ class TopLevelMixin(mixins.ConfigMethodMixin):
     def to_dict(
         self,
         validate: bool = True,
+        *,
+        format: str = "vega-lite",
         ignore: Optional[List[str]] = None,
         context: Optional[TypingDict[str, Any]] = None,
     ) -> dict:
@@ -829,6 +832,8 @@ class TopLevelMixin(mixins.ConfigMethodMixin):
         validate : bool, optional
             If True (default), then validate the output dictionary
             against the schema.
+        format : str, optional
+            Chart specification format, one of "vega-lite" (default) or "vega"
         ignore : list[str], optional
             A list of keys to ignore. It is usually not needed
             to specify this argument as a user.
@@ -851,6 +856,13 @@ class TopLevelMixin(mixins.ConfigMethodMixin):
         SchemaValidationError
             if validate=True and the dict does not conform to the schema
         """
+
+        # Validate format
+        if format not in ("vega-lite", "vega"):
+            raise ValueError(
+                f'The format argument must be either "vega-lite" or "vega". Received {repr(format)}'
+            )
+
         # We make use of three context markers:
         # - 'data' points to the data that should be referenced for column type
         #   inference.
@@ -880,28 +892,78 @@ class TopLevelMixin(mixins.ConfigMethodMixin):
         # TopLevelMixin instance does not necessarily have to_dict defined
         # but due to how Altair is set up this should hold.
         # Too complex to type hint right now
-        dct = super(TopLevelMixin, copy).to_dict(  # type: ignore[misc]
+        vegalite_spec = super(TopLevelMixin, copy).to_dict(  # type: ignore[misc]
             validate=validate, ignore=ignore, context=context
         )
 
         # TODO: following entries are added after validation. Should they be validated?
         if is_top_level:
             # since this is top-level we add $schema if it's missing
-            if "$schema" not in dct:
-                dct["$schema"] = SCHEMA_URL
+            if "$schema" not in vegalite_spec:
+                vegalite_spec["$schema"] = SCHEMA_URL
 
             # apply theme from theme registry
             the_theme = themes.get()
             # Use assert to tell type checkers that it is not None. Holds true
             # as there is always a default theme set when importing Altair
             assert the_theme is not None
-            dct = utils.update_nested(the_theme(), dct, copy=True)
+            vegalite_spec = utils.update_nested(the_theme(), vegalite_spec, copy=True)
 
             # update datasets
             if context["datasets"]:
-                dct.setdefault("datasets", {}).update(context["datasets"])
+                vegalite_spec.setdefault("datasets", {}).update(context["datasets"])
+
+        if format == "vega":
+            plugin = vegalite_compilers.get()
+            if plugin is None:
+                raise ValueError("No active vega-lite compiler plugin found")
+            return plugin(vegalite_spec)
+        else:
+            return vegalite_spec
+
+    def to_json(
+        self,
+        validate: bool = True,
+        indent: int = 2,
+        sort_keys: bool = True,
+        *,
+        format: str = "vega-lite",
+        ignore: Optional[List[str]] = None,
+        context: Optional[TypingDict[str, Any]] = None,
+        **kwargs,
+    ) -> str:
+        """Convert a chart to a JSON string
 
-        return dct
+        Parameters
+        ----------
+        validate : bool, optional
+            If True (default), then validate the output dictionary
+            against the schema.
+        indent : int, optional
+            The number of spaces of indentation to use. The default is 2.
+        sort_keys : bool, optional
+            If True (default), sort keys in the output.
+        format : str, optional
+            The chart specification format. One of "vega-lite" (default) or "vega".
+            The "vega" format relies on the active Vega-Lite compiler plugin, which
+            by default requires the vl-convert-python package.
+        ignore : list[str], optional
+            A list of keys to ignore. It is usually not needed
+            to specify this argument as a user.
+        context : dict[str, Any], optional
+            A context dictionary. It is usually not needed
+            to specify this argument as a user.
+        **kwargs
+            Additional keyword arguments are passed to ``json.dumps()``
+        """
+        if ignore is None:
+            ignore = []
+        if context is None:
+            context = {}
+        spec = self.to_dict(
+            validate=validate, format=format, ignore=ignore, context=context
+        )
+        return json.dumps(spec, indent=indent, sort_keys=sort_keys, **kwargs)
 
     def to_html(
         self,
@@ -2546,6 +2608,8 @@ class Chart(
     def to_dict(
         self,
         validate: bool = True,
+        *,
+        format: str = "vega-lite",
         ignore: Optional[List[str]] = None,
         context: Optional[TypingDict[str, Any]] = None,
     ) -> dict:
@@ -2556,6 +2620,8 @@ class Chart(
         validate : bool, optional
             If True (default), then validate the output dictionary
             against the schema.
+        format : str, optional
+            Chart specification format, one of "vega-lite" (default) or "vega"
         ignore : list[str], optional
             A list of keys to ignore. It is usually not needed
             to specify this argument as a user.
@@ -2585,9 +2651,11 @@ class Chart(
             copy = self.copy(deep=False)
             copy.data = core.InlineData(values=[{}])
             return super(Chart, copy).to_dict(
-                validate=validate, ignore=ignore, context=context
+                validate=validate, format=format, ignore=ignore, context=context
             )
-        return super().to_dict(validate=validate, ignore=ignore, context=context)
+        return super().to_dict(
+            validate=validate, format=format, ignore=ignore, context=context
+        )
 
     def add_params(self, *params) -> Self:
         """Add one or more parameters to the chart."""
diff --git a/altair/vegalite/v5/compiler.py b/altair/vegalite/v5/compiler.py
new file mode 100644
index 00000000..60afa99f
--- /dev/null
+++ b/altair/vegalite/v5/compiler.py
@@ -0,0 +1,27 @@
+from ...utils.compiler import VegaLiteCompilerRegistry
+
+ENTRY_POINT_GROUP: str = "altair.vegalite.v5.vegalite_compiler"
+vegalite_compilers = VegaLiteCompilerRegistry(entry_point_group=ENTRY_POINT_GROUP)
+
+
+def vl_convert_compiler(vegalite_spec: dict) -> dict:
+    """
+    Vega-Lite to Vega compiler that uses vl-convert
+    """
+    from . import SCHEMA_VERSION
+
+    try:
+        import vl_convert as vlc
+    except ImportError as err:
+        raise ImportError(
+            "The vl-convert Vega-Lite compiler requires the vl-convert-python package"
+        ) from err
+
+    # Compute vl-convert's vl_version string (of the form 'v5_8')
+    # from SCHEMA_VERSION (of the form 'v5.8.0')
+    vl_version = "_".join(SCHEMA_VERSION.split(".")[:2])
+    return vlc.vegalite_to_vega(vegalite_spec, vl_version=vl_version)
+
+
+vegalite_compilers.register("vl-convert", vl_convert_compiler)
+vegalite_compilers.enable("vl-convert")
diff --git a/tools/schemapi/schemapi.py b/tools/schemapi/schemapi.py
index 7d1a73a1..7809acaf 100644
--- a/tools/schemapi/schemapi.py
+++ b/tools/schemapi/schemapi.py
@@ -735,6 +735,7 @@ class SchemaBase:
     def to_dict(
         self,
         validate: bool = True,
+        *,
         ignore: Optional[List[str]] = None,
         context: Optional[Dict[str, Any]] = None,
     ) -> dict:
@@ -827,6 +828,7 @@ class SchemaBase:
         validate: bool = True,
         indent: int = 2,
         sort_keys: bool = True,
+        *,
         ignore: Optional[List[str]] = None,
         context: Optional[Dict[str, Any]] = None,
         **kwargs,


Test patch:
diff --git a/tests/utils/test_compiler.py b/tests/utils/test_compiler.py
new file mode 100644
index 00000000..7805dc08
--- /dev/null
+++ b/tests/utils/test_compiler.py
@@ -0,0 +1,41 @@
+import json
+import pytest
+from altair import vegalite_compilers, Chart
+
+
+@pytest.fixture
+def chart():
+    return (
+        Chart("cars.json")
+        .mark_point()
+        .encode(
+            x="Horsepower:Q",
+            y="Miles_per_Gallon:Q",
+        )
+    )
+
+
+def assert_is_vega_spec(vega_spec):
+    assert vega_spec["$schema"] == "https://vega.github.io/schema/vega/v5.json"
+    assert "data" in vega_spec
+    assert "marks" in vega_spec
+    assert "scales" in vega_spec
+    assert "axes" in vega_spec
+
+
+def test_vegalite_compiler(chart):
+    vegalite_spec = chart.to_dict()
+    vega_spec = vegalite_compilers.get()(vegalite_spec)
+    assert_is_vega_spec(vega_spec)
+
+
+def test_to_dict_with_format_vega(chart):
+    vega_spec = chart.to_dict(format="vega")
+    assert_is_vega_spec(vega_spec)
+
+
+def test_to_json_with_format_vega(chart):
+    json_spec = chart.to_json(format="vega")
+    assert isinstance(json_spec, str)
+    spec = json.loads(json_spec)
+    assert_is_vega_spec(spec)