Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 60 additions & 19 deletions ggsql-python/python/ggsql/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@

from ggsql._ggsql import (
DuckDBReader,
VegaLiteWriter,
VegaLiteWriter as _RustVegaLiteWriter,
Validated,
Spec,
validate,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)
39 changes: 39 additions & 0 deletions ggsql-python/tests/test_ggsql.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Loading