Skip to content

Commit f2a8bd3

Browse files
cpsievertclaude
andcommitted
feat(python): add VegaLiteWriter.render_chart() for Altair output
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent f5d0f85 commit f2a8bd3

2 files changed

Lines changed: 107 additions & 19 deletions

File tree

ggsql-python/python/ggsql/__init__.py

Lines changed: 60 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010

1111
from ggsql._ggsql import (
1212
DuckDBReader,
13-
VegaLiteWriter,
13+
VegaLiteWriter as _RustVegaLiteWriter,
1414
Validated,
1515
Spec,
1616
validate,
@@ -75,6 +75,64 @@ def register(
7575
) -> None: ...
7676

7777

78+
def _json_to_altair_chart(vegalite_json: str, **kwargs: Any) -> AltairChart:
79+
"""Convert a Vega-Lite JSON string to the appropriate Altair chart type."""
80+
spec = json.loads(vegalite_json)
81+
82+
if "layer" in spec:
83+
return altair.LayerChart.from_json(vegalite_json, **kwargs)
84+
elif "facet" in spec or "spec" in spec:
85+
return altair.FacetChart.from_json(vegalite_json, **kwargs)
86+
elif "concat" in spec:
87+
return altair.ConcatChart.from_json(vegalite_json, **kwargs)
88+
elif "hconcat" in spec:
89+
return altair.HConcatChart.from_json(vegalite_json, **kwargs)
90+
elif "vconcat" in spec:
91+
return altair.VConcatChart.from_json(vegalite_json, **kwargs)
92+
elif "repeat" in spec:
93+
return altair.RepeatChart.from_json(vegalite_json, **kwargs)
94+
else:
95+
return altair.Chart.from_json(vegalite_json, **kwargs)
96+
97+
98+
class VegaLiteWriter:
99+
"""Vega-Lite v6 JSON output writer.
100+
101+
Methods
102+
-------
103+
render(spec)
104+
Render a Spec to a Vega-Lite JSON string.
105+
render_chart(spec, **kwargs)
106+
Render a Spec to an Altair chart object.
107+
"""
108+
109+
def __init__(self) -> None:
110+
self._inner = _RustVegaLiteWriter()
111+
112+
def render(self, spec: Spec) -> str:
113+
"""Render a Spec to a Vega-Lite JSON string."""
114+
return self._inner.render(spec)
115+
116+
def render_chart(self, spec: Spec, **kwargs: Any) -> AltairChart:
117+
"""Render a Spec to an Altair chart object.
118+
119+
Parameters
120+
----------
121+
spec
122+
The resolved visualization specification from ``reader.execute()``.
123+
**kwargs
124+
Additional keyword arguments passed to ``altair.Chart.from_json()``.
125+
Common options include ``validate=False`` to skip schema validation.
126+
127+
Returns
128+
-------
129+
AltairChart
130+
An Altair chart object (Chart, LayerChart, FacetChart, etc.).
131+
"""
132+
vegalite_json = self.render(spec)
133+
return _json_to_altair_chart(vegalite_json, **kwargs)
134+
135+
78136
def render_altair(
79137
df: IntoFrame,
80138
viz: str,
@@ -120,21 +178,4 @@ def render_altair(
120178
writer = VegaLiteWriter()
121179
vegalite_json = writer.render(spec)
122180

123-
# Parse to determine the correct Altair class
124-
spec = json.loads(vegalite_json)
125-
126-
# Determine the correct Altair class based on spec structure
127-
if "layer" in spec:
128-
return altair.LayerChart.from_json(vegalite_json, **kwargs)
129-
elif "facet" in spec or "spec" in spec:
130-
return altair.FacetChart.from_json(vegalite_json, **kwargs)
131-
elif "concat" in spec:
132-
return altair.ConcatChart.from_json(vegalite_json, **kwargs)
133-
elif "hconcat" in spec:
134-
return altair.HConcatChart.from_json(vegalite_json, **kwargs)
135-
elif "vconcat" in spec:
136-
return altair.VConcatChart.from_json(vegalite_json, **kwargs)
137-
elif "repeat" in spec:
138-
return altair.RepeatChart.from_json(vegalite_json, **kwargs)
139-
else:
140-
return altair.Chart.from_json(vegalite_json, **kwargs)
181+
return _json_to_altair_chart(vegalite_json, **kwargs)

ggsql-python/tests/test_ggsql.py

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -618,6 +618,53 @@ def test_reader_is_exported(self):
618618
assert hasattr(ggsql, "Reader")
619619

620620

621+
class TestVegaLiteWriterRenderChart:
622+
"""Tests for VegaLiteWriter.render_chart() method."""
623+
624+
def test_render_chart_returns_altair_chart(self):
625+
"""render_chart() returns an Altair chart object."""
626+
reader = ggsql.DuckDBReader("duckdb://memory")
627+
spec = reader.execute("SELECT 1 AS x, 2 AS y VISUALISE x, y DRAW point")
628+
writer = ggsql.VegaLiteWriter()
629+
chart = writer.render_chart(spec)
630+
assert isinstance(chart, altair.TopLevelMixin)
631+
632+
def test_render_chart_layer(self):
633+
"""render_chart() returns LayerChart for layered specs."""
634+
reader = ggsql.DuckDBReader("duckdb://memory")
635+
spec = reader.execute("SELECT 1 AS x, 2 AS y VISUALISE x, y DRAW point")
636+
writer = ggsql.VegaLiteWriter()
637+
chart = writer.render_chart(spec)
638+
assert isinstance(chart, altair.LayerChart)
639+
640+
def test_render_chart_facet(self):
641+
"""render_chart() returns FacetChart for faceted specs."""
642+
reader = ggsql.DuckDBReader("duckdb://memory")
643+
df = pl.DataFrame(
644+
{
645+
"x": [1, 2, 3, 4, 5, 6],
646+
"y": [10, 20, 30, 40, 50, 60],
647+
"group": ["A", "A", "A", "B", "B", "B"],
648+
}
649+
)
650+
reader.register("data", df)
651+
spec = reader.execute(
652+
"SELECT * FROM data VISUALISE x, y FACET group DRAW point"
653+
)
654+
writer = ggsql.VegaLiteWriter()
655+
chart = writer.render_chart(spec, validate=False)
656+
assert isinstance(chart, altair.FacetChart)
657+
658+
def test_render_chart_kwargs_forwarded(self):
659+
"""render_chart() forwards kwargs to from_json()."""
660+
reader = ggsql.DuckDBReader("duckdb://memory")
661+
spec = reader.execute("SELECT 1 AS x, 2 AS y VISUALISE x, y DRAW point")
662+
writer = ggsql.VegaLiteWriter()
663+
# Should not raise (validate=False is forwarded)
664+
chart = writer.render_chart(spec, validate=False)
665+
assert isinstance(chart, altair.TopLevelMixin)
666+
667+
621668
class TestTypeStubs:
622669
"""Tests for type stub presence and correctness."""
623670

0 commit comments

Comments
 (0)