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}