Skip to main content

utils/
shell_expander.rs

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