kroki_rs/diagrams/providers/d2.rs
1use crate::diagrams::{DiagramError, DiagramProvider, DiagramResult};
2use async_trait::async_trait;
3use std::process::Stdio;
4use tokio::process::Command;
5
6crate::diagrams::define_provider!(D2Provider);
7
8#[async_trait]
9impl DiagramProvider for D2Provider {
10 fn validate(&self, source: &str) -> DiagramResult<()> {
11 if source.trim().is_empty() {
12 return Err(DiagramError::ValidationFailed(
13 "Diagram source is empty".into(),
14 ));
15 }
16 Ok(())
17 }
18
19 async fn generate(&self, source: &str, format: &str) -> DiagramResult<Vec<u8>> {
20 // d2 - - reads from stdin and writes to stdout
21 // But only if format is svg?
22 // d2 supports --stdout-format json|svg|png|...
23
24 // Let's check d2 help again for --stdout-format default.
25 // It says "d2 compiles ... to file.svg ... defaults to file.svg".
26 // "Use - to have d2 read from stdin or write to stdout."
27
28 // If I use `d2 - -`, it writes SVG to stdout (default).
29 // If I want other formats, I need to specify?
30 // `d2 input.d2 output.png`
31 // `d2 --stdout-format png - -` ?
32
33 let mut cmd = Command::new(&self.bin_path);
34
35 // Input is stdin: "-"
36 // Output is stdout: "-"
37
38 // Need to handle formats.
39 // d2 supports: svg, png, pdf, pptx, gif, txt
40
41 // If format is passed, we might need a flag.
42 // d2 usage: d2 [flags] input output
43
44 // If format is svg (default): `d2 - -` works.
45 // If format is png: `d2 input.d2 output.png`.
46 // Does `d2 - -` support changing format?
47 // Help says: `--stdout-format string output format when writing to stdout ... Usage: d2 input.d2 --stdout-format png - > output.png`
48
49 // So correct usage for stdout is: `d2 --stdout-format <format> - -` (input -, output - implicit or explicit?)
50 // The help example `d2 input.d2 --stdout-format png -` has input file `input.d2` and output `-`.
51
52 // So if input is `-`, use `d2 --stdout-format <format> - -`?
53 // Let's assume input is `-`.
54
55 match format {
56 "svg" | "png" | "pdf" => {
57 // OK
58 }
59 _ => {
60 return Err(DiagramError::UnsupportedFormat {
61 format: format.into(),
62 provider: "D2".into(),
63 })
64 }
65 }
66
67 cmd.arg("--layout=dagre"); // Default layout, maybe make configurable?
68 // Actually don't enforce layout unless needed.
69
70 // Use --stdout-format
71 cmd.arg("--stdout-format").arg(format);
72
73 // Input file: "-" (stdin)
74 cmd.arg("-");
75
76 // Output file: "-" (stdout) - Wait, if --stdout-format is used, maybe output file argument is not needed or must be `-`?
77 // Help example: `d2 input.d2 --stdout-format png -`
78 // Here `input.d2` is arg1, `-` is arg2 (output).
79 // So if input is `-`: `d2 --stdout-format png - -`
80 cmd.arg("-");
81
82 cmd.stdin(Stdio::piped())
83 .stdout(Stdio::piped())
84 .stderr(Stdio::piped());
85
86 let output = crate::diagrams::run_process_with_timeout(
87 "d2",
88 cmd,
89 Some(source.as_bytes()),
90 self.timeout_ms,
91 source.len(),
92 )
93 .await?;
94
95 if output.status.success() {
96 if output.stdout.is_empty() {
97 return Err(DiagramError::ProcessFailed(
98 "D2 conversion succeeded but output is empty".into(),
99 ));
100 }
101 Ok(output.stdout)
102 } else {
103 let stderr = String::from_utf8_lossy(&output.stderr);
104 Err(DiagramError::ProcessFailed(format!(
105 "D2 conversion failed: {}",
106 stderr
107 )))
108 }
109 }
110}