utils/
variables.rs

1use regex::{Captures, Regex};
2use std::borrow::Cow;
3use std::collections::BTreeMap;
4use std::env;
5use std::fmt;
6use std::sync::{Arc, LazyLock};
7use thiserror::Error;
8
9type GetEnvVarFn = Arc<dyn Fn(&str) -> Option<String> + Send + Sync>;
10
11/// A small key→value lookup used when expanding `$VAR` / `${VAR}` references.
12///
13/// `Vars` is checked before falling back to an env lookup (the process
14/// environment by default), so callers can inject locally-scoped values
15/// (e.g. `WORKSPACE`) without leaking them into global state.
16#[derive(Clone)]
17pub struct Vars {
18    map: BTreeMap<String, String>,
19    get_env_var: GetEnvVarFn,
20}
21
22impl fmt::Debug for Vars {
23    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
24        f.debug_struct("Vars").field("map", &self.map).finish_non_exhaustive()
25    }
26}
27
28impl Default for Vars {
29    fn default() -> Self {
30        Self::new()
31    }
32}
33
34impl Vars {
35    pub fn new() -> Self {
36        Self { map: BTreeMap::new(), get_env_var: Arc::new(|k| env::var(k).ok()) }
37    }
38
39    /// Override the env-var fallback. Useful in tests to avoid mutating the
40    /// process environment.
41    pub fn with_env_lookup(mut self, f: impl Fn(&str) -> Option<String> + Send + Sync + 'static) -> Self {
42        self.get_env_var = Arc::new(f);
43        self
44    }
45
46    pub fn with(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
47        self.insert(key, value);
48        self
49    }
50
51    pub fn insert(&mut self, key: impl Into<String>, value: impl Into<String>) -> &mut Self {
52        self.map.insert(key.into(), value.into());
53        self
54    }
55
56    pub fn get(&self, key: &str) -> Option<Cow<'_, str>> {
57        if let Some(value) = self.map.get(key) {
58            Some(Cow::Borrowed(value.as_str()))
59        } else {
60            (self.get_env_var)(key).map(Cow::Owned)
61        }
62    }
63
64    /// Expand `$VAR` / `${VAR}` references in `template`, consulting `self` before
65    /// the env lookup. `$$` escapes to a literal `$`.
66    pub fn expand(&self, template: &str) -> Result<String, VarError> {
67        let mut missing = None;
68        let expanded = REFERENCE.replace_all(template, |caps: &Captures| {
69            let Some(name) = var_name(caps) else {
70                return "$".to_string();
71            };
72            self.get(name).map_or_else(
73                || {
74                    missing.get_or_insert_with(|| name.to_string());
75                    caps[0].to_string()
76                },
77                Cow::into_owned,
78            )
79        });
80        match missing {
81            Some(var) => Err(VarError::NotFound(var)),
82            None => Ok(expanded.into_owned()),
83        }
84    }
85
86    /// Returns `true` if `s` contains a `$VAR` or `${VAR}` style reference.
87    pub fn has_reference(s: &str) -> bool {
88        REFERENCE.captures_iter(s).any(|caps| var_name(&caps).is_some())
89    }
90}
91
92#[derive(Debug, Error)]
93pub enum VarError {
94    #[error("Environment variable '{0}' not found")]
95    NotFound(String),
96}
97
98static REFERENCE: LazyLock<Regex> =
99    LazyLock::new(|| Regex::new(r"\$(?:\$|\{([^}]+)\}|([A-Za-z_][A-Za-z0-9_]*))").expect("valid regex"));
100
101fn var_name<'a>(caps: &Captures<'a>) -> Option<&'a str> {
102    caps.get(1).or_else(|| caps.get(2)).map(|name| name.as_str())
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108
109    fn no_env() -> Vars {
110        Vars::new().with_env_lookup(|_| None)
111    }
112
113    #[test]
114    fn test_simple_var() {
115        let vars = no_env().with("TEST_VAR_SIMPLE", "hello");
116        let result = vars.expand("$TEST_VAR_SIMPLE world").unwrap();
117        assert_eq!(result, "hello world");
118    }
119
120    #[test]
121    fn test_bracketed_var() {
122        let vars = no_env().with("TEST_VAR_BRACKET", "hello");
123        let result = vars.expand("${TEST_VAR_BRACKET} world").unwrap();
124        assert_eq!(result, "hello world");
125    }
126
127    #[test]
128    fn test_escape_sequence() {
129        let result = no_env().expand("$$VAR").unwrap();
130        assert_eq!(result, "$VAR");
131    }
132
133    #[test]
134    fn test_multiple_vars() {
135        let vars = no_env().with("VAR1", "hello").with("VAR2", "world");
136        let result = vars.expand("$VAR1 ${VAR2}!").unwrap();
137        assert_eq!(result, "hello world!");
138    }
139
140    #[test]
141    fn test_missing_var() {
142        let result = no_env().expand("$MISSING_VAR");
143        assert!(matches!(result, Err(VarError::NotFound(ref name)) if name == "MISSING_VAR"));
144    }
145
146    #[test]
147    fn test_unclosed_brace_left_as_is() {
148        let result = no_env().expand("${VAR").unwrap();
149        assert_eq!(result, "${VAR");
150    }
151
152    #[test]
153    fn test_empty_string() {
154        let result = no_env().expand("").unwrap();
155        assert_eq!(result, "");
156    }
157
158    #[test]
159    fn test_no_vars() {
160        let result = no_env().expand("plain text").unwrap();
161        assert_eq!(result, "plain text");
162    }
163
164    #[test]
165    fn test_dollar_at_end() {
166        let result = no_env().expand("text$").unwrap();
167        assert_eq!(result, "text$");
168    }
169
170    #[test]
171    fn test_var_with_underscore() {
172        let vars = no_env().with("MY_TEST_VAR", "value");
173        let result = vars.expand("$MY_TEST_VAR").unwrap();
174        assert_eq!(result, "value");
175    }
176
177    #[test]
178    fn test_var_with_numbers() {
179        let vars = no_env().with("VAR123", "value");
180        let result = vars.expand("$VAR123").unwrap();
181        assert_eq!(result, "value");
182    }
183
184    #[test]
185    fn test_special_char_stops_var_name() {
186        let vars = no_env().with("VAR", "value");
187        let result = vars.expand("$VAR-suffix").unwrap();
188        assert_eq!(result, "value-suffix");
189    }
190
191    #[test]
192    fn vars_lookup_takes_precedence_over_env() {
193        let vars = Vars::new()
194            .with_env_lookup(|k| (k == "WORKSPACE").then(|| "/from-env".into()))
195            .with("WORKSPACE", "/from-vars");
196        let result = vars.expand("${WORKSPACE}/foo").unwrap();
197        assert_eq!(result, "/from-vars/foo");
198    }
199
200    #[test]
201    fn vars_falls_through_to_env_when_missing() {
202        let vars = Vars::new().with_env_lookup(|k| (k == "ONLY_IN_ENV").then(|| "from-env".into()));
203        let result = vars.expand("$ONLY_IN_ENV").unwrap();
204        assert_eq!(result, "from-env");
205    }
206
207    #[test]
208    fn substituted_values_are_not_expanded_again() {
209        let vars = no_env().with("A", "$B").with("B", "expanded twice");
210        assert_eq!(vars.expand("${A} $A").unwrap(), "$B $B");
211    }
212
213    #[test]
214    fn has_reference_ignores_escaped_dollars() {
215        assert!(!Vars::has_reference("$$HOME"));
216    }
217
218    #[test]
219    fn has_reference_detects_both_forms() {
220        assert!(Vars::has_reference("${WORKSPACE}/foo"));
221        assert!(Vars::has_reference("$HOME/bar"));
222        assert!(!Vars::has_reference("plain/path"));
223        assert!(!Vars::has_reference("$"));
224        assert!(!Vars::has_reference("$$"));
225    }
226}
utils/variables.rs

utils/
variables.rs
