vegafusion_core/spec/
chart.rs

1use crate::error::{Result, ResultWithContext, VegaFusionError};
2use crate::proto::gen::tasks::{Task, TzConfig};
3use crate::spec::axis::AxisSpec;
4use crate::spec::data::DataSpec;
5use crate::spec::mark::MarkSpec;
6use crate::spec::projection::ProjectionSpec;
7use crate::spec::scale::ScaleSpec;
8use crate::spec::signal::SignalSpec;
9use crate::spec::title::TitleSpec;
10use crate::spec::visitors::{
11    DefinitionVarsChartVisitor, InputVarsChartVisitor, MakeTaskScopeVisitor, MakeTasksVisitor,
12    UpdateVarsChartVisitor,
13};
14use crate::task_graph::graph::ScopedVariable;
15use crate::task_graph::scope::TaskScope;
16use itertools::sorted;
17use serde::{Deserialize, Serialize};
18use serde_json::Value;
19use std::collections::HashMap;
20
21#[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)]
22pub struct ChartSpec {
23    #[serde(rename = "$schema", default = "default_schema")]
24    pub schema: String,
25
26    #[serde(default, skip_serializing_if = "Vec::is_empty")]
27    pub data: Vec<DataSpec>,
28
29    #[serde(default, skip_serializing_if = "Vec::is_empty")]
30    pub signals: Vec<SignalSpec>,
31
32    #[serde(default, skip_serializing_if = "Vec::is_empty")]
33    pub marks: Vec<MarkSpec>,
34
35    #[serde(default, skip_serializing_if = "Vec::is_empty")]
36    pub scales: Vec<ScaleSpec>,
37
38    #[serde(default, skip_serializing_if = "Vec::is_empty")]
39    pub projections: Vec<ProjectionSpec>,
40
41    #[serde(default, skip_serializing_if = "Vec::is_empty")]
42    pub axes: Vec<AxisSpec>,
43
44    #[serde(skip_serializing_if = "Option::is_none")]
45    pub title: Option<TitleSpec>,
46
47    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
48    pub usermeta: HashMap<String, Value>,
49
50    #[serde(flatten)]
51    pub extra: HashMap<String, Value>,
52}
53
54pub fn default_schema() -> String {
55    String::from("https://vega.github.io/schema/vega/v5.json")
56}
57
58impl ChartSpec {
59    pub fn walk(&self, visitor: &mut dyn ChartVisitor) -> Result<()> {
60        // Child groups
61        let scope: Vec<u32> = Vec::new();
62        let mut group_index = 0;
63        for mark in &self.marks {
64            if mark.type_ == "group" {
65                // Add group index to scope
66                let mut nested_scope = scope.clone();
67                nested_scope.push(group_index);
68
69                visitor.visit_group_mark(mark, &nested_scope)?;
70                mark.walk(visitor, &nested_scope)?;
71                group_index += 1;
72            } else {
73                // Keep parent scope
74                visitor.visit_non_group_mark(mark, &scope)?;
75            }
76        }
77
78        // Top-level with empty scope
79        for data in &self.data {
80            visitor.visit_data(data, &scope)?;
81        }
82        for scale in &self.scales {
83            visitor.visit_scale(scale, &scope)?;
84        }
85        for projection in &self.projections {
86            visitor.visit_projection(projection, &scope)?;
87        }
88        for axis in &self.axes {
89            visitor.visit_axis(axis, &scope)?;
90        }
91        for signal in &self.signals {
92            visitor.visit_signal(signal, &scope)?;
93        }
94
95        // Visit top-level chart
96        visitor.visit_chart(self)?;
97
98        Ok(())
99    }
100
101    pub fn walk_mut(&mut self, visitor: &mut dyn MutChartVisitor) -> Result<()> {
102        // Visit top-level chart
103        visitor.visit_chart(self)?;
104
105        // Top-level with empty scope
106        let scope: Vec<u32> = Vec::new();
107        for data in &mut self.data {
108            visitor.visit_data(data, &scope)?;
109        }
110        for scale in &mut self.scales {
111            visitor.visit_scale(scale, &scope)?;
112        }
113        for projection in &mut self.projections {
114            visitor.visit_projection(projection, &scope)?;
115        }
116        for axis in &mut self.axes {
117            visitor.visit_axis(axis, &scope)?;
118        }
119        for signal in &mut self.signals {
120            visitor.visit_signal(signal, &scope)?;
121        }
122
123        // Child groups
124        let mut group_index = 0;
125        for mark in &mut self.marks {
126            if mark.type_ == "group" {
127                // Add group index to scope
128                let mut nested_scope = scope.clone();
129                nested_scope.push(group_index);
130
131                visitor.visit_group_mark(mark, &nested_scope)?;
132                mark.walk_mut(visitor, &nested_scope)?;
133                group_index += 1;
134            } else {
135                // Keep parent scope
136                visitor.visit_non_group_mark(mark, &scope)?;
137            }
138        }
139
140        Ok(())
141    }
142
143    pub fn to_task_scope(&self) -> Result<TaskScope> {
144        let mut visitor = MakeTaskScopeVisitor::new();
145        self.walk(&mut visitor)?;
146        Ok(visitor.task_scope)
147    }
148
149    pub fn to_tasks(
150        &self,
151        tz_config: &TzConfig,
152        dataset_fingerprints: &HashMap<String, String>,
153    ) -> Result<Vec<Task>> {
154        let mut visitor = MakeTasksVisitor::new(tz_config, dataset_fingerprints);
155        self.walk(&mut visitor)?;
156        Ok(visitor.tasks)
157    }
158
159    pub fn get_group(&self, group_index: u32) -> Result<&MarkSpec> {
160        self.marks
161            .iter()
162            .filter(|m| m.type_ == "group")
163            .nth(group_index as usize)
164            .with_context(|| format!("No group with index {group_index}"))
165    }
166
167    pub fn get_nested_group(&self, path: &[u32]) -> Result<&MarkSpec> {
168        if path.is_empty() {
169            return Err(VegaFusionError::internal(
170                "Nested group scope may not be empty",
171            ));
172        }
173        let mut group = self.get_group(path[0])?;
174        for group_index in &path[1..] {
175            group = group.get_group(*group_index)?;
176        }
177        Ok(group)
178    }
179
180    pub fn get_group_mut(&mut self, group_index: u32) -> Result<&mut MarkSpec> {
181        self.marks
182            .iter_mut()
183            .filter(|m| m.type_ == "group")
184            .nth(group_index as usize)
185            .with_context(|| format!("No group with index {group_index}"))
186    }
187
188    pub fn get_nested_group_mut(&mut self, path: &[u32]) -> Result<&mut MarkSpec> {
189        if path.is_empty() {
190            return Err(VegaFusionError::internal("Path may not be empty"));
191        }
192        let mut group = self.get_group_mut(path[0])?;
193        for group_index in &path[1..] {
194            group = group.get_group_mut(*group_index)?;
195        }
196        Ok(group)
197    }
198
199    pub fn get_nested_signal(&self, path: &[u32], name: &str) -> Result<&SignalSpec> {
200        let signals = if path.is_empty() {
201            &self.signals
202        } else {
203            let group = self.get_nested_group(path)?;
204            &group.signals
205        };
206        signals
207            .iter()
208            .find(|s| s.name == name)
209            .with_context(|| format!("No signal named {name} found at path {path:?}"))
210    }
211
212    pub fn get_nested_signal_mut(&mut self, path: &[u32], name: &str) -> Result<&mut SignalSpec> {
213        let signals = if path.is_empty() {
214            &mut self.signals
215        } else {
216            let group = self.get_nested_group_mut(path)?;
217            &mut group.signals
218        };
219        signals
220            .iter_mut()
221            .find(|s| s.name == name)
222            .with_context(|| format!("No signal named {name} found at path {path:?}"))
223    }
224
225    pub fn get_nested_data(&self, path: &[u32], name: &str) -> Result<&DataSpec> {
226        let datasets = if path.is_empty() {
227            &self.data
228        } else {
229            let group = self.get_nested_group(path)?;
230            &group.data
231        };
232        datasets
233            .iter()
234            .find(|s| s.name == name)
235            .with_context(|| format!("No data named {name} found at path {path:?}"))
236    }
237
238    pub fn get_nested_data_mut(&mut self, path: &[u32], name: &str) -> Result<&mut DataSpec> {
239        let signals = if path.is_empty() {
240            &mut self.data
241        } else {
242            let group = self.get_nested_group_mut(path)?;
243            &mut group.data
244        };
245        signals
246            .iter_mut()
247            .find(|s| s.name == name)
248            .with_context(|| format!("No data named {name} found at path {path:?}"))
249    }
250
251    pub fn add_nested_signal(
252        &mut self,
253        path: &[u32],
254        spec: SignalSpec,
255        index: Option<usize>,
256    ) -> Result<()> {
257        let signals = if path.is_empty() {
258            &mut self.signals
259        } else {
260            let group = self.get_nested_group_mut(path)?;
261            &mut group.signals
262        };
263        match index {
264            Some(index) => {
265                signals.insert(index, spec);
266            }
267            None => {
268                signals.push(spec);
269            }
270        }
271        Ok(())
272    }
273
274    pub fn add_nested_data(
275        &mut self,
276        path: &[u32],
277        spec: DataSpec,
278        index: Option<usize>,
279    ) -> Result<()> {
280        let data = if path.is_empty() {
281            &mut self.data
282        } else {
283            let group = self.get_nested_group_mut(path)?;
284            &mut group.data
285        };
286        match index {
287            Some(index) => {
288                data.insert(index, spec);
289            }
290            None => {
291                data.push(spec);
292            }
293        }
294        Ok(())
295    }
296
297    pub fn definition_vars(&self) -> Result<Vec<ScopedVariable>> {
298        let mut visitor = DefinitionVarsChartVisitor::new();
299        self.walk(&mut visitor)?;
300        Ok(sorted(visitor.definition_vars).collect())
301    }
302
303    pub fn update_vars(&self, task_scope: &TaskScope) -> Result<Vec<ScopedVariable>> {
304        let mut visitor = UpdateVarsChartVisitor::new(task_scope);
305        self.walk(&mut visitor)?;
306        Ok(sorted(visitor.update_vars).collect())
307    }
308
309    pub fn input_vars(&self, task_scope: &TaskScope) -> Result<Vec<ScopedVariable>> {
310        let mut visitor = InputVarsChartVisitor::new(task_scope);
311        self.walk(&mut visitor)?;
312        Ok(sorted(visitor.input_vars).collect())
313    }
314}
315
316pub trait ChartVisitor {
317    fn visit_chart(&mut self, _chart: &ChartSpec) -> Result<()> {
318        Ok(())
319    }
320    fn visit_data(&mut self, _data: &DataSpec, _scope: &[u32]) -> Result<()> {
321        Ok(())
322    }
323    fn visit_signal(&mut self, _signal: &SignalSpec, _scope: &[u32]) -> Result<()> {
324        Ok(())
325    }
326    fn visit_scale(&mut self, _scale: &ScaleSpec, _scope: &[u32]) -> Result<()> {
327        Ok(())
328    }
329    fn visit_projection(&mut self, _projection: &ProjectionSpec, _scope: &[u32]) -> Result<()> {
330        Ok(())
331    }
332    fn visit_axis(&mut self, _axis: &AxisSpec, _scope: &[u32]) -> Result<()> {
333        Ok(())
334    }
335    fn visit_non_group_mark(&mut self, _mark: &MarkSpec, _scope: &[u32]) -> Result<()> {
336        Ok(())
337    }
338    fn visit_group_mark(&mut self, _mark: &MarkSpec, _scope: &[u32]) -> Result<()> {
339        Ok(())
340    }
341}
342
343pub trait MutChartVisitor {
344    fn visit_chart(&mut self, _chart: &mut ChartSpec) -> Result<()> {
345        Ok(())
346    }
347    fn visit_data(&mut self, _data: &mut DataSpec, _scope: &[u32]) -> Result<()> {
348        Ok(())
349    }
350    fn visit_signal(&mut self, _signal: &mut SignalSpec, _scope: &[u32]) -> Result<()> {
351        Ok(())
352    }
353    fn visit_scale(&mut self, _scale: &mut ScaleSpec, _scope: &[u32]) -> Result<()> {
354        Ok(())
355    }
356    fn visit_projection(&mut self, _projection: &mut ProjectionSpec, _scope: &[u32]) -> Result<()> {
357        Ok(())
358    }
359    fn visit_axis(&mut self, _axis: &mut AxisSpec, _scope: &[u32]) -> Result<()> {
360        Ok(())
361    }
362    fn visit_non_group_mark(&mut self, _mark: &mut MarkSpec, _scope: &[u32]) -> Result<()> {
363        Ok(())
364    }
365    fn visit_group_mark(&mut self, _mark: &mut MarkSpec, _scope: &[u32]) -> Result<()> {
366        Ok(())
367    }
368}
vegafusion_core/spec/chart.rs

vegafusion_core/spec/
chart.rs
