Skip to content
Merged
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
135 changes: 134 additions & 1 deletion app/charts.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,89 @@
# Every data-source name a spec may reference; anything else in a `data` node is dropped.
_ALLOWED_SOURCES = frozenset({_RESULT_SOURCE, *_GEO_ASSETS})

# Vega's built-in color scheme names (https://vega.github.io/vega/docs/schemes/). A scheme
# name outside this set is rejected. Includes the dark-/light-background ramp variants for
# completeness, though the prompt steers away from them. Keep in sync with the scheme list
# in _CHART_SPEC_INSTRUCTIONS.
_VEGA_COLOR_SCHEMES = frozenset(
{
# Categorical
"accent",
"category10",
"category20",
"category20b",
"category20c",
"observable10",
"dark2",
"paired",
"pastel1",
"pastel2",
"set1",
"set2",
"set3",
"tableau10",
"tableau20",
# Sequential single-hue
"blues",
"tealblues",
"teals",
"greens",
"browns",
"oranges",
"reds",
"purples",
"warmgreys",
"greys",
# Sequential multi-hue
"viridis",
"magma",
"inferno",
"plasma",
"cividis",
"turbo",
"bluegreen",
"bluepurple",
"goldgreen",
"goldorange",
"goldred",
"greenblue",
"orangered",
"purplebluegreen",
"purpleblue",
"purplered",
"redpurple",
"yellowgreenblue",
"yellowgreen",
"yelloworangebrown",
"yelloworangered",
# Sequential multi-hue for dark / light backgrounds
"darkblue",
"darkgold",
"darkgreen",
"darkmulti",
"darkred",
"lightgreyred",
"lightgreyteal",
"lightmulti",
"lightorange",
"lighttealblue",
# Diverging
"blueorange",
"brownbluegreen",
"purplegreen",
"pinkyellowgreen",
"purpleorange",
"redblue",
"redgrey",
"redyellowblue",
"redyellowgreen",
"spectral",
# Cyclical
"rainbow",
"sinebow",
}
)

# System prompt for the chart spec generation.
_CHART_SPEC_INSTRUCTIONS = """\
You are a data visualization specialist. Given a small, already-aggregated query result and a description of the chart to build, return one complete Vega-Lite v6 spec. Reference the result's columns by their exact names, and never include a data source, dataset, or URL — the exact rows are injected separately. You may use composite views and transforms when they make the chart clearer.
Expand All @@ -62,7 +145,19 @@

## Color

Never set any color-related property — no `scale`, `scheme`, `sort`, `range`, or explicit color value — on any encoding, mark, or config, whether the field is categorical, quantitative, or there's only one series. Just map the field to the color channel and leave everything else about its color to the frontend's defaults.
By default, set no color-related property — no `scale`, `scheme`, `sort`, `range`, or explicit color value — on any encoding, mark, or config, whether the field is categorical, quantitative, or there's only one series. Just map the field to the color channel and leave everything else about its color to the frontend's brand palette.

The one exception is when the description explicitly asks for a specific color. Then honor exactly what was asked, choosing the encoding that fits the chart and the number of series being colored, and set only that — leave every other color to the frontend. Do not invent a color, scheme, or palette the description did not ask for.

Give the color as explicit CSS colors (a hex string or a standard color keyword) whenever the request names a color rather than a palette: a single fixed color as a constant color `value`, and specific per-category colors as a `range` on the color encoding's `scale`. For a gradient over a quantitative or ordinal field, give a `range` of two or more explicit colors, e.g. light-to-dark for low-to-high.

Use a named `scheme` only when the request names or clearly maps to one of the built-in Vega schemes below, and then use its name exactly. Never guess or invent a scheme name. If no built-in scheme matches the requested color, use an explicit `range` instead.

- Categorical: `accent`, `category10`, `category20`, `category20b`, `category20c`, `observable10`, `dark2`, `paired`, `pastel1`, `pastel2`, `set1`, `set2`, `set3`, `tableau10`, `tableau20`.
- Sequential single-hue: `blues`, `tealblues`, `teals`, `greens`, `browns`, `oranges`, `reds`, `purples`, `warmgreys`, `greys`.
- Sequential multi-hue: `viridis`, `magma`, `inferno`, `plasma`, `cividis`, `turbo`, `bluegreen`, `bluepurple`, `goldgreen`, `goldorange`, `goldred`, `greenblue`, `orangered`, `purplebluegreen`, `purpleblue`, `purplered`, `redpurple`, `yellowgreenblue`, `yellowgreen`, `yelloworangebrown`, `yelloworangered`.
- Diverging: `blueorange`, `brownbluegreen`, `purplegreen`, `pinkyellowgreen`, `purpleorange`, `redblue`, `redgrey`, `redyellowblue`, `redyellowgreen`, `spectral`.
- Cyclical: `rainbow`, `sinebow`.

## Number formatting

Expand Down Expand Up @@ -480,6 +575,34 @@ def _collect(node: JsonValue, key: str) -> set[str]:
return found


def _collect_scheme_names(node: JsonValue) -> set[str]:
"""Collect every color-scheme name a spec references, anywhere within it.

A `scheme` is either a string (`"scheme": "reds"`) or a `SchemeParams` object whose
`name` holds the scheme (`"scheme": {"name": "reds", "extent": [...]}`); this collects
the name from both forms so `_validate_chart_spec` can reject an unknown one.

Args:
node (JsonValue): A spec, or any node within it, to walk.

Returns:
set[str]: Every scheme name found under a `scheme` key.
"""
found: set[str] = set()
if isinstance(node, dict):
for key, value in node.items():
if key == "scheme":
if isinstance(value, str):
found.add(value)
elif isinstance(value, dict) and isinstance(value.get("name"), str):
found.add(value["name"])
found |= _collect_scheme_names(value)
elif isinstance(node, list):
for item in node:
found |= _collect_scheme_names(item)
return found


def _validate_chart_spec(spec: dict[str, Any], columns: list[str]) -> list[str]:
"""Return the reasons a spec would not render a correct chart (empty if valid).

Expand All @@ -505,6 +628,16 @@ def _validate_chart_spec(spec: dict[str, Any], columns: list[str]) -> list[str]:
f"Available columns: {sorted(columns)}."
)

# A `scheme` name outside Vega's built-in set errors only once real data instantiates the
# scale — which vl_convert swallows below (it logs to stderr but still returns an SVG), so
# the render check never catches it and the chart ships blank. Reject unknown names here.
unknown_schemes = sorted(_collect_scheme_names(spec) - _VEGA_COLOR_SCHEMES)
if unknown_schemes:
errors.append(
f"Unknown color scheme(s): {unknown_schemes}. Use an exact Vega scheme name, "
"or set explicit colors via a `range` instead."
)

# Resolve named sources (stub geometry, empty rows) so a geoshape/lookup spec renders
# its geometry; a normal chart just gets empty data. Rendering (not just VL→Vega compiling)
# is the only way to validate expression strings in the spec.
Expand Down
136 changes: 133 additions & 3 deletions tests/app/test_charts.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
ChartSpecInvalid,
_chart_spec_user_prompt,
_collect,
_collect_scheme_names,
_fetch_rows,
_geo_stub_node,
_geo_url_node,
Expand Down Expand Up @@ -334,6 +335,50 @@ def test_ignores_non_string_values_under_the_key(self):
assert _collect({"field": 5, "x": {"field": "ano"}}, "field") == {"ano"}


class TestCollectSchemeNames:
def test_collects_a_string_scheme(self):
spec = {"encoding": {"color": {"field": "v", "scale": {"scheme": "reds"}}}}

assert _collect_scheme_names(spec) == {"reds"}

def test_collects_a_scheme_params_object_by_name(self):
# `scale.scheme` may be a SchemeParams object; the name lives under `name`.
spec = {
"encoding": {
"color": {
"field": "v",
"scale": {"scheme": {"name": "viridis", "extent": [0.2, 1]}},
}
}
}

assert _collect_scheme_names(spec) == {"viridis"}

def test_collects_every_scheme_across_layers(self):
spec = {
"layer": [
{"encoding": {"color": {"scale": {"scheme": "reds"}}}},
{"encoding": {"color": {"scale": {"scheme": {"name": "blues"}}}}},
]
}

assert _collect_scheme_names(spec) == {"reds", "blues"}

def test_returns_empty_when_no_scheme(self):
spec = {
"mark": "bar",
"encoding": {"color": {"scale": {"range": ["red", "blue"]}}},
}

assert _collect_scheme_names(spec) == set()

def test_ignores_a_scheme_object_without_a_string_name(self):
# A malformed scheme object (no usable name) yields nothing, not a crash.
spec = {"encoding": {"color": {"scale": {"scheme": {"extent": [0, 1]}}}}}

assert _collect_scheme_names(spec) == set()


class TestValidateChartSpec:
"""Uses the real vl-convert compiler (in-process)."""

Expand Down Expand Up @@ -390,8 +435,8 @@ def test_invalid_mark_fails_to_compile(self):

assert any("compile" in error for error in errors)

def test_unknown_color_scheme_fails_to_compile(self):
"""An invalid scheme (d3's RdBu) is rejected by the compile step, like any bad value."""
def test_unknown_color_scheme_is_rejected(self):
"""A scheme name outside Vega's built-in set (here d3's `RdBu`) is flagged."""
spec = {
"mark": "rect",
"encoding": {
Expand All @@ -407,7 +452,92 @@ def test_unknown_color_scheme_fails_to_compile(self):

errors = _validate_chart_spec(spec, ["ano", "mes", "temp"])

assert any("compile" in error for error in errors)
assert any("RdBu" in error and "scheme" in error for error in errors)

def test_unknown_scheme_in_choropleth_is_caught_by_the_allowlist(self):
"""The regression that shipped a blank chart.

A geoshape's bad scheme does not raise in vl-convert (it logs to stderr but still
returns an SVG), so the compile step alone never catches it — the allowlist does.
"""
spec = _choropleth_spec()
spec["encoding"]["color"]["scale"] = {"scheme": "lightpinkdark"}

errors = _validate_chart_spec(spec, ["sigla_uf", "valor"])

assert any("lightpinkdark" in error for error in errors)
# The compile step swallows this one, so it is not what caught it.
assert not any("compile" in error for error in errors)

def test_unknown_scheme_params_object_is_rejected(self):
"""A bad scheme given as a SchemeParams object is caught too."""
spec = {
"mark": "bar",
"encoding": {
"x": {"field": "ano", "type": "ordinal"},
"y": {"field": "total", "type": "quantitative"},
"color": {
"field": "total",
"type": "quantitative",
"scale": {"scheme": {"name": "notarealscheme"}},
},
},
}

errors = _validate_chart_spec(spec, ["ano", "total"])

assert any("notarealscheme" in error for error in errors)

def test_valid_scheme_passes(self):
"""A real Vega scheme (`reds`) is accepted."""
spec = {
"mark": "rect",
"encoding": {
"x": {"field": "ano", "type": "ordinal"},
"y": {"field": "mes", "type": "ordinal"},
"color": {
"field": "temp",
"type": "quantitative",
"scale": {"scheme": "reds"},
},
},
}

assert _validate_chart_spec(spec, ["ano", "mes", "temp"]) == []

def test_valid_scheme_params_object_passes(self):
"""A real scheme given as a SchemeParams object is accepted."""
spec = {
"mark": "bar",
"encoding": {
"x": {"field": "ano", "type": "ordinal"},
"y": {"field": "total", "type": "quantitative"},
"color": {
"field": "total",
"type": "quantitative",
"scale": {"scheme": {"name": "viridis"}},
},
},
}

assert _validate_chart_spec(spec, ["ano", "total"]) == []

def test_explicit_color_range_needs_no_scheme(self):
"""Explicit `range` colors (no scheme) pass the scheme check."""
spec = {
"mark": "bar",
"encoding": {
"x": {"field": "ano", "type": "nominal"},
"y": {"field": "total", "type": "quantitative"},
"color": {
"field": "ano",
"type": "nominal",
"scale": {"domain": ["2023", "2024"], "range": ["red", "blue"]},
},
},
}

assert _validate_chart_spec(spec, ["ano", "total"]) == []

def test_choropleth_spec_is_valid(self):
"""A choropleth compiles (real geometry) and its geo properties are allowed."""
Expand Down
Loading