Skip to main content

utils/
shell_expander.rs

1use futures::future::join_all;
2use regex::Regex;
3use std::fmt::{Display, Formatter};
4use std::path::Path;
5use std::{env, fmt};
6use tokio::process::Command;
7
8/// Expands `` !`command` `` markers in text by running each command via
9/// `$SHELL -c` (fallback `sh`) and substituting the trimmed stdout.
10///
11/// Construct once and reuse across a batch of expansions to amortize regex
12/// compilation.
13pub struct ShellExpander {
14    regex: Regex,
15}
16
17impl ShellExpander {
18    pub fn new() -> Self {
19        Self { regex: Regex::new(r"!`([^`\n]+)`").expect("valid regex") }
20    }
21
22    /// Expand `` !`command` `` markers in `content`, running each command from
23    /// `cwd`. Returns `content` unchanged if no markers are present.
24    ///
25    /// Markers are expanded concurrently; the first non-zero exit or spawn
26    /// failure short-circuits and surfaces as [`ShellInterpError`].
27    pub async fn expand(&self, content: &str, cwd: &Path) -> String {
28        if !self.regex.is_match(content) {
29            return content.to_string();
30        }
31
32        let spans: Vec<(usize, usize, &str)> = self
33            .regex
34            .captures_iter(content)
35            .filter_map(|captures| {
36                let whole = captures.get(0)?;
37                let cmd = captures.get(1)?;
38                Some((whole.start(), whole.end(), cmd.as_str()))
39            })
40            .collect();
41
42        let outputs = join_all(spans.iter().map(|(_, _, cmd)| Self::run(cmd, cwd))).await;
43        let mut out = String::with_capacity(content.len());
44        let mut last = 0;
45
46        for ((start, end, _), result) in spans.iter().zip(outputs.into_iter()) {
47            out.push_str(&content[last..*start]);
48            match result {
49                Ok(output) => out.push_str(&output),
50                Err(err) => tracing::warn!("{err}"),
51            }
52            last = *end;
53        }
54
55        out.push_str(&content[last..]);
56        out
57    }
58
59    async fn run(cmd: &str, cwd: &Path) -> Result<String, ShellExpansionError> {
60        let shell = env::var("SHELL").unwrap_or_else(|_| "sh".to_string());
61        let output = Command::new(&shell).arg("-c").arg(cmd).current_dir(cwd).output().await.map_err(|e| {
62            ShellExpansionError::Spawn { shell: shell.clone(), cmd: cmd.to_string(), error: e.to_string() }
63        })?;
64
65        if !output.status.success() {
66            let stderr = String::from_utf8_lossy(&output.stderr);
67            return Err(ShellExpansionError::NonZeroExit {
68                cmd: cmd.to_string(),
69                status: output.status.to_string(),
70                stderr: stderr.trim().to_string(),
71            });
72        }
73
74        Ok(String::from_utf8_lossy(&output.stdout).trim_end().to_string())
75    }
76}
77
78#[derive(Debug)]
79pub enum ShellExpansionError {
80    Spawn { shell: String, cmd: String, error: String },
81    NonZeroExit { cmd: String, status: String, stderr: String },
82}
83
84impl Default for ShellExpander {
85    fn default() -> Self {
86        Self::new()
87    }
88}
89
90impl Display for ShellExpansionError {
91    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
92        match self {
93            Self::Spawn { shell, cmd, error } => {
94                write!(f, "Failed to spawn {shell} for `{cmd}`: {error}")
95            }
96            Self::NonZeroExit { cmd, status, stderr } => {
97                write!(f, "Shell interpolation `{cmd}` failed with {status}: {stderr}")
98            }
99        }
100    }
101}
102
103impl std::error::Error for ShellExpansionError {}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108
109    #[tokio::test]
110    async fn no_op_without_marker() {
111        let content = "Just some plain content with no directives";
112        let expander = ShellExpander::new();
113        let cwd = std::env::current_dir().unwrap();
114        let result = expander.expand(content, &cwd).await;
115        assert_eq!(result, content);
116    }
117
118    #[tokio::test]
119    async fn runs_shell_command() {
120        let expander = ShellExpander::new();
121        let cwd = std::env::current_dir().unwrap();
122        let result = expander.expand("branch: !`echo main`", &cwd).await;
123        assert_eq!(result, "branch: main");
124    }
125
126    #[tokio::test]
127    async fn runs_command_in_cwd() {
128        let dir = tempfile::tempdir().unwrap();
129        std::fs::write(dir.path().join("sentinel.txt"), "").unwrap();
130
131        let expander = ShellExpander::new();
132        let result = expander.expand("files: !`ls`", dir.path()).await;
133        assert!(result.contains("sentinel.txt"), "expected sentinel.txt in output: {result}");
134    }
135
136    #[tokio::test]
137    async fn handles_multiple_commands() {
138        let expander = ShellExpander::new();
139        let cwd = std::env::current_dir().unwrap();
140        let result = expander.expand("a=!`echo one`, b=!`echo two`", &cwd).await;
141        assert_eq!(result, "a=one, b=two");
142    }
143
144    #[tokio::test]
145    async fn failed_command_substitutes_empty_string() {
146        let expander = ShellExpander::new();
147        let cwd = std::env::current_dir().unwrap();
148        let result = expander.expand("before !`exit 1` after", &cwd).await;
149        assert_eq!(result, "before  after");
150    }
151
152    #[tokio::test]
153    async fn trims_trailing_whitespace() {
154        let expander = ShellExpander::new();
155        let cwd = std::env::current_dir().unwrap();
156        let result = expander.expand("!`printf 'hi\\n\\n'`", &cwd).await;
157        assert_eq!(result, "hi");
158    }
159}