From ccbe5bbe31e87773bcb592d3e01f19cf236bb440 Mon Sep 17 00:00:00 2001 From: Carson Date: Mon, 2 Mar 2026 16:15:49 -0600 Subject: [PATCH] feat(python): add VegaLiteWriter.render_chart() for Altair output Co-Authored-By: Claude Opus 4.6 --- ggsql-python/python/ggsql/__init__.py | 79 ++++++++++++++++++++------- ggsql-python/tests/test_ggsql.py | 39 +++++++++++++ 2 files changed, 99 insertions(+), 19 deletions(-) diff --git a/ggsql-python/python/ggsql/__init__.py b/ggsql-python/python/ggsql/__init__.py index d69c84ef..f9356e72 100644 --- a/ggsql-python/python/ggsql/__init__.py +++ b/ggsql-python/python/ggsql/__init__.py @@ -9,7 +9,7 @@ from ggsql._ggsql import ( DuckDBReader, - VegaLiteWriter, + VegaLiteWriter as _RustVegaLiteWriter, Validated, Spec, validate, @@ -41,6 +41,64 @@ ] +def _json_to_altair_chart(vegalite_json: str, **kwargs: Any) -> AltairChart: + """Convert a Vega-Lite JSON string to the appropriate Altair chart type.""" + spec = json.loads(vegalite_json) + + if "layer" in spec: + return altair.LayerChart.from_json(vegalite_json, **kwargs) + elif "facet" in spec or "spec" in spec: + return altair.FacetChart.from_json(vegalite_json, **kwargs) + elif "concat" in spec: + return altair.ConcatChart.from_json(vegalite_json, **kwargs) + elif "hconcat" in spec: + return altair.HConcatChart.from_json(vegalite_json, **kwargs) + elif "vconcat" in spec: + return altair.VConcatChart.from_json(vegalite_json, **kwargs) + elif "repeat" in spec: + return altair.RepeatChart.from_json(vegalite_json, **kwargs) + else: + return altair.Chart.from_json(vegalite_json, **kwargs) + + +class VegaLiteWriter: + """Vega-Lite v6 JSON output writer. + + Methods + ------- + render(spec) + Render a Spec to a Vega-Lite JSON string. + render_chart(spec, **kwargs) + Render a Spec to an Altair chart object. + """ + + def __init__(self) -> None: + self._inner = _RustVegaLiteWriter() + + def render(self, spec: Spec) -> str: + """Render a Spec to a Vega-Lite JSON string.""" + return self._inner.render(spec) + + def render_chart(self, spec: Spec, **kwargs: Any) -> AltairChart: + """Render a Spec to an Altair chart object. + + Parameters + ---------- + spec + The resolved visualization specification from ``reader.execute()``. + **kwargs + Additional keyword arguments passed to ``altair.Chart.from_json()``. + Common options include ``validate=False`` to skip schema validation. + + Returns + ------- + AltairChart + An Altair chart object (Chart, LayerChart, FacetChart, etc.). + """ + vegalite_json = self.render(spec) + return _json_to_altair_chart(vegalite_json, **kwargs) + + def render_altair( df: IntoFrame, viz: str, @@ -86,21 +144,4 @@ def render_altair( writer = VegaLiteWriter() vegalite_json = writer.render(spec) - # Parse to determine the correct Altair class - spec = json.loads(vegalite_json) - - # Determine the correct Altair class based on spec structure - if "layer" in spec: - return altair.LayerChart.from_json(vegalite_json, **kwargs) - elif "facet" in spec or "spec" in spec: - return altair.FacetChart.from_json(vegalite_json, **kwargs) - elif "concat" in spec: - return altair.ConcatChart.from_json(vegalite_json, **kwargs) - elif "hconcat" in spec: - return altair.HConcatChart.from_json(vegalite_json, **kwargs) - elif "vconcat" in spec: - return altair.VConcatChart.from_json(vegalite_json, **kwargs) - elif "repeat" in spec: - return altair.RepeatChart.from_json(vegalite_json, **kwargs) - else: - return altair.Chart.from_json(vegalite_json, **kwargs) + return _json_to_altair_chart(vegalite_json, **kwargs) diff --git a/ggsql-python/tests/test_ggsql.py b/ggsql-python/tests/test_ggsql.py index 8b12f103..fbe4b131 100644 --- a/ggsql-python/tests/test_ggsql.py +++ b/ggsql-python/tests/test_ggsql.py @@ -530,3 +530,42 @@ def unregister(self, name: str) -> None: writer = ggsql.VegaLiteWriter() json_output = writer.render(spec) assert "point" in json_output + + +class TestVegaLiteWriterRenderChart: + """Tests for VegaLiteWriter.render_chart() method.""" + + def test_render_chart_returns_altair_chart(self): + """render_chart() returns an Altair chart object.""" + reader = ggsql.DuckDBReader("duckdb://memory") + spec = reader.execute("SELECT 1 AS x, 2 AS y VISUALISE x, y DRAW point") + writer = ggsql.VegaLiteWriter() + chart = writer.render_chart(spec) + assert isinstance(chart, altair.TopLevelMixin) + + def test_render_chart_layer(self): + """render_chart() returns LayerChart for layered specs.""" + reader = ggsql.DuckDBReader("duckdb://memory") + spec = reader.execute("SELECT 1 AS x, 2 AS y VISUALISE x, y DRAW point") + writer = ggsql.VegaLiteWriter() + chart = writer.render_chart(spec) + assert isinstance(chart, altair.LayerChart) + + def test_render_chart_facet(self): + """render_chart() returns FacetChart for faceted specs.""" + reader = ggsql.DuckDBReader("duckdb://memory") + df = pl.DataFrame( + { + "x": [1, 2, 3, 4, 5, 6], + "y": [10, 20, 30, 40, 50, 60], + "group": ["A", "A", "A", "B", "B", "B"], + } + ) + reader.register("data", df) + spec = reader.execute( + "SELECT * FROM data VISUALISE x, y FACET group DRAW point" + ) + writer = ggsql.VegaLiteWriter() + chart = writer.render_chart(spec, validate=False) + assert isinstance(chart, altair.FacetChart) +