kroki_rs/diagrams/providers/
vega.rs

1use crate::diagrams::{DiagramError, DiagramProvider, DiagramResult};
2use async_trait::async_trait;
3use std::path::PathBuf;
4use std::process::Stdio;
5use tokio::process::Command;
6
7crate::diagrams::define_provider!(VegaProvider);
8
9#[async_trait]
10impl DiagramProvider for VegaProvider {
11    fn validate(&self, source: &str) -> DiagramResult<()> {
12        if source.trim().is_empty() {
13            return Err(DiagramError::ValidationFailed(
14                "Diagram source is empty".into(),
15            ));
16        }
17        Ok(())
18    }
19
20    async fn generate(&self, source: &str, format: &str) -> DiagramResult<Vec<u8>> {
21        if format != "svg" {
22            return Err(DiagramError::UnsupportedFormat {
23                format: format.into(),
24                provider: "Vega".into(),
25            });
26        }
27
28        let mut cmd = Command::new(&self.bin_path);
29        cmd.stdin(Stdio::piped())
30            .stdout(Stdio::piped())
31            .stderr(Stdio::piped());
32
33        let output = crate::diagrams::run_process_with_timeout(
34            "vg2svg",
35            cmd,
36            Some(source.as_bytes()),
37            self.timeout_ms,
38            source.len(),
39        )
40        .await?;
41
42        if output.status.success() {
43            if output.stdout.is_empty() {
44                return Err(DiagramError::ProcessFailed(
45                    "Vega conversion succeeded but output is empty".into(),
46                ));
47            }
48            Ok(output.stdout)
49        } else {
50            let stderr = String::from_utf8_lossy(&output.stderr);
51            Err(DiagramError::ProcessFailed(format!(
52                "Vega conversion failed: {}",
53                stderr
54            )))
55        }
56    }
57}
58
59pub struct VegaLiteProvider {
60    pub vl_bin_path: PathBuf, // vl2vg
61    pub vg_bin_path: PathBuf, // vg2svg
62    pub timeout_ms: Option<u64>,
63}
64
65impl VegaLiteProvider {
66    pub fn new(vl_bin_path: PathBuf, vg_bin_path: PathBuf, timeout_ms: Option<u64>) -> Self {
67        Self {
68            vl_bin_path,
69            vg_bin_path,
70            timeout_ms,
71        }
72    }
73}
74
75#[async_trait]
76impl DiagramProvider for VegaLiteProvider {
77    fn validate(&self, source: &str) -> DiagramResult<()> {
78        if source.trim().is_empty() {
79            return Err(DiagramError::ValidationFailed(
80                "Diagram source is empty".into(),
81            ));
82        }
83        Ok(())
84    }
85
86    async fn generate(&self, source: &str, format: &str) -> DiagramResult<Vec<u8>> {
87        if format != "svg" {
88            return Err(DiagramError::UnsupportedFormat {
89                format: format.into(),
90                provider: "Vega-Lite".into(),
91            });
92        }
93
94        // Stage 1: vl2vg (Vega-Lite spec → Vega spec)
95        let mut vl_cmd = Command::new(&self.vl_bin_path);
96        vl_cmd
97            .stdin(Stdio::piped())
98            .stdout(Stdio::piped())
99            .stderr(Stdio::piped());
100
101        let vl_output = crate::diagrams::run_process_with_timeout(
102            "vl2vg",
103            vl_cmd,
104            Some(source.as_bytes()),
105            self.timeout_ms,
106            source.len(),
107        )
108        .await?;
109
110        if !vl_output.status.success() {
111            let stderr = String::from_utf8_lossy(&vl_output.stderr);
112            return Err(DiagramError::ProcessFailed(format!(
113                "vl2vg (stage 1) failed: {}",
114                stderr
115            )));
116        }
117
118        // Stage 2: vg2svg (Vega spec → SVG)
119        let vg_input = &vl_output.stdout;
120        let mut vg_cmd = Command::new(&self.vg_bin_path);
121        vg_cmd
122            .stdin(Stdio::piped())
123            .stdout(Stdio::piped())
124            .stderr(Stdio::piped());
125
126        let vg_output = crate::diagrams::run_process_with_timeout(
127            "vg2svg",
128            vg_cmd,
129            Some(vg_input.as_slice()),
130            self.timeout_ms,
131            vg_input.len(),
132        )
133        .await?;
134
135        if vg_output.status.success() {
136            if vg_output.stdout.is_empty() {
137                return Err(DiagramError::ProcessFailed(format!(
138                    "vg2svg (stage 2) succeeded but output is empty (vl2vg produced {} bytes)",
139                    vg_input.len()
140                )));
141            }
142            Ok(vg_output.stdout)
143        } else {
144            let stderr = String::from_utf8_lossy(&vg_output.stderr);
145            Err(DiagramError::ProcessFailed(format!(
146                "vg2svg (stage 2) failed (vl2vg produced {} bytes): {}",
147                vg_input.len(),
148                stderr
149            )))
150        }
151    }
152}