recursive_env/

lib.rs

use anyhow::{bail, Context, Result};
use std::env;
// TODO: we dont really need anyhow and can drop the dep
// TODO: should we have used a parser combinator library?

// TODO: We don't need to handle escape characters, right? It wouldn't make sense because they
// aren't valid for var names
// TODO: shoule we also implement our own equlivalent of std::env::var_os(), std::env::vars(), and std::env::set_var()?
// TODO: we should have a sub function that evaluates without doing the acutal check of the env
// vars. This way we can test that function without setting actual env vars and then run the tests
// in parallel again
// TODO: stop evaluation of variables with out braces on double quotes
// TODO: allow escaping of '$' ('\$' and '\${}')
// TOOD: it would be really cool if the test cases ran a bash shell and matched the output with ops
// done in bash or maybe just sh

/// Read a variable from the environment and look up an other env variables contained within its
/// definition.
pub fn var(key: &str) -> Result<String> {
    let value = env::var(key).with_context(|| {
        format!(
            "'{}' not found in env. Require to evaulate another var",
            key,
        )
    })?;

    // TODO: lots of nested logic here. Maybe we need a state machine or something. Or a
    // parser combinator

    let mut previous_frame: Option<String> = None;
    let mut current_frame = String::new();
    let mut stop_on_whitespace = false;
    let mut in_double_quotes = false;
    let mut capturing = false;
    let mut add_next_quote = false;
    let mut single_quote_no_special_handeling = false;

    let mut chars = value.chars().peekable();
    loop {
        let Some(current_char) = chars.next() else {
            break;
        };

        if single_quote_no_special_handeling {
            if current_char == '\\' {
                // Add escaped backslash to the current frame
                let Some(next_char) = chars.next() else {
                    bail!("Unable to parse: Lone '\\' at end of input");
                };
                match next_char {
                    '"' => {
                        current_frame.push(next_char);
                    }
                    _ => {
                        // Keep the escape char
                        current_frame.push(current_char);
                        current_frame.push(next_char);
                    }
                }
                continue;
            }

            // TODO: need to handle escaped single quote
            current_frame.push(current_char);
        } else {
            match current_char {
                '\\' => {
                    let Some(next_char) = chars.next() else {
                        bail!("Unable to parse: Lone '\\' at end of input");
                    };

                    match next_char {
                        '"' => {
                            current_frame.push(next_char);
                        }
                        _ => {
                            // Keep the escape char
                            current_frame.push(current_char);
                            current_frame.push(next_char);
                        }
                    }

                    continue;
                }
                '$' => {
                    let Some(next_char) = chars.peek() else {
                        bail!("Unable to parse: Lone '$' at end of input");
                    };

                    // Pause the current frame and start a new one
                    previous_frame = Some(current_frame);
                    current_frame = String::new();

                    match next_char {
                        &'{' => {
                            // Actually move on to next char to skip the '{' (was peeked previously)
                            chars.next().unwrap();
                        }
                        &'(' => {
                            // "$(...)" messes up our check for standalone '$' without a following '{'
                            // Push the dollar sign and move on as usual. We aren't going to try and
                            // evaulate a sub expression
                            current_frame.push(current_char);
                        }
                        _ => {
                            stop_on_whitespace = true;
                            capturing = true;
                        }
                    }
                    continue;
                }
                '}' => {
                    // We've finished pushing a variable name to the current frame. Look it up
                    let value = var(&current_frame).with_context(|| {
                        format!(
                            "'{}' not found in env. Require to evaulate another var",
                            value,
                        )
                    })?;
                    // Jump back to the last frame and append our value in place of where the '${KEY}'
                    // would have been
                    let Some(ref prev) = previous_frame else {
                        bail!("Found an unmatched '}}'");
                    };
                    current_frame = prev.to_owned();
                    current_frame.push_str(&value);
                    continue;
                }
                ' ' | '\t' => {
                    if stop_on_whitespace {
                        // TODO: most of this code is duplicated with the match '}' branch
                        // We've finished pushing a variable name to the current frame. Look it up
                        let value = var(&current_frame).with_context(|| {
                            format!(
                                "'{}' not found in env. Require to evaulate another var",
                                value,
                            )
                        })?;
                        // Jump back to the last frame and append our value in place of where the '${KEY}'
                        // would have been
                        let Some(ref prev) = previous_frame else {
                            bail!("Error. TODO: better error message");
                        };
                        current_frame = prev.to_owned();
                        current_frame.push_str(&value);
                        // Now that we've done all that we can push the whitespace on
                        stop_on_whitespace = false;
                        capturing = false;
                        // we will continue on to add the whitespace to the current frame
                    }
                }
                '\'' => {
                    // Toggle single_quote_no_special_handeling
                    single_quote_no_special_handeling = !single_quote_no_special_handeling;
                }
                '"' => {
                    if in_double_quotes {
                        if add_next_quote {
                            current_frame.push(current_char);
                            add_next_quote = false;
                        }

                        // We've found the closing quote
                        if capturing {
                            // TODO: most of this code is duplicated with the match '}' branch
                            // We've finished pushing a variable name to the current frame. Look it up
                            let value = var(&current_frame).with_context(|| {
                                format!(
                                    "'{}' not found in env. Require to evaulate another var",
                                    value,
                                )
                            })?;
                            // Jump back to the last frame and append our value in place of where the '${KEY}'
                            // would have been
                            let Some(ref prev) = previous_frame else {
                                bail!("Error. TODO: better error message");
                            };
                            current_frame = prev.to_owned();
                            current_frame.push_str(&value);
                            capturing = false;
                        }
                    } else {
                        in_double_quotes = true;

                        let Some(next_char) = chars.peek() else {
                            bail!("Unable to parse: Lone '\"' at end of input");
                        };

                        if next_char != &'$' {
                            add_next_quote = true; // Denote that we need to capture the closing
                                                   // quote
                            current_frame.push(current_char);
                        }
                    }

                    continue;
                }
                _ => {
                    // No need to do anything special for any other char
                }
            }
            current_frame.push(current_char);
        }
    }

    if capturing {
        // Finish out the final var we were working on
        // TODO: This is all duplicate code

        // We've finished pushing a variable name to the current frame. Look it up
        let value = var(&current_frame).with_context(|| {
            format!(
                "'{}' not found in env. Require to evaulate another var",
                value,
            )
        })?;
        // Jump back to the last frame and append our value in place of where the '${KEY}'
        // would have been
        let Some(ref prev) = previous_frame else {
            bail!("Error. TODO: better error message");
        };
        current_frame = prev.to_owned();
        current_frame.push_str(&value);
    }

    Ok(current_frame)
}

#[cfg(test)]
mod tests {
    use super::*;
    use pretty_assertions::{assert_eq, assert_ne};
    use serial_test::serial;

    // I think that when running in parallel the setting/unsetting of env vars was causing the
    // tests to mess with eachother

    #[test]
    #[serial]
    fn test_nonrecursive() {
        let key = "KEY";
        let value = "VALUE".to_string();
        env::set_var(key, &value);
        assert_eq!(var(key).unwrap(), value);
    }

    #[test]
    #[serial]
    fn test_recursive() {
        let key1 = "KEY1";
        let key2 = "KEY2";
        env::set_var(key2, "number2");
        env::set_var(key1, "number1 ${KEY2} number3");
        assert_eq!(var(key2).unwrap(), "number2".to_string());
        assert_eq!(var(key1).unwrap(), "number1 number2 number3".to_string());
    }

    #[test]
    #[serial]
    fn test_more_recursive() {
        let key1 = "KEY1";
        let key2 = "KEY2";
        let key3 = "KEY3";
        env::set_var(key3, "number3");
        env::set_var(key2, "number2 ${KEY3} number4");
        env::set_var(key1, "number1 ${KEY2} number5");
        assert_eq!(var(key3).unwrap(), "number3".to_string());
        assert_eq!(var(key2).unwrap(), "number2 number3 number4".to_string());
        assert_eq!(
            var(key1).unwrap(),
            "number1 number2 number3 number4 number5".to_string()
        );
    }

    #[test]
    #[serial]
    fn test_key_not_found() {
        let key = "KEY";
        let _ = env::remove_var(key); // Ignore an error if it comes up. This probably means the
                                      // key wasn't set to begin with
        assert!(var(key).is_err());
    }

    #[test]
    #[serial]
    fn test_real_world_example() {
        env::set_var("HOME", "/home/agaia");
        env::set_var("XDG_DATA_HOME", "${HOME}/.local/share");
        env::set_var("DATABASE_URL", "sqlite:${XDG_DATA_HOME}/taskrs/data.db");
        let db_path = var("DATABASE_URL").unwrap();
        assert_eq!(db_path, "sqlite:/home/agaia/.local/share/taskrs/data.db");
    }

    #[test]
    #[serial]
    fn test_recursive_without_braces() {
        let key1 = "KEY1";
        let key2 = "KEY2";
        env::set_var(key2, "number2");
        env::set_var(key1, "number1 $KEY2 number3");
        assert_eq!(var(key2).unwrap(), "number2".to_string());
        assert_eq!(var(key1).unwrap(), "number1 number2 number3".to_string());
    }

    #[test]
    #[serial]
    fn test_more_recursive_without_braces() {
        let key1 = "KEY1";
        let key2 = "KEY2";
        let key3 = "KEY3";
        env::set_var(key3, "number3");
        env::set_var(key2, "number2 $KEY3 number4");
        env::set_var(key1, "number1 $KEY2 number5");
        assert_eq!(var(key3).unwrap(), "number3".to_string());
        assert_eq!(var(key2).unwrap(), "number2 number3 number4".to_string());
        assert_eq!(
            var(key1).unwrap(),
            "number1 number2 number3 number4 number5".to_string()
        );
    }

    #[test]
    #[serial]
    fn test_no_braces_stop_on_whitespace() {
        let key1 = "KEY1";
        let key1_longer = "KEY1BUTLONGER";
        let key2 = "KEY2";
        env::set_var(key1, "1");
        env::set_var(key1_longer, "2");
        env::set_var(key2, "prefix$KEY1BUTLONGER ");
        assert_eq!(var(key2).unwrap(), "prefix2 ".to_string());
    }

    #[test]
    #[serial]
    fn test_no_braces_no_ending_whitespace() {
        let key = "KEY";
        let key2 = "KEY2";
        env::set_var(key, "test");
        env::set_var(key2, "$KEY");
        assert_eq!(var(key2).unwrap(), "test".to_string());
    }

    #[test]
    #[serial]
    fn test_do_not_eval_subexpression() {
        let key = "KEY";
        let value = String::from("$(subexpression)");
        env::set_var(&key, &value);
        assert_eq!(var(key).unwrap(), value);
    }

    #[test]
    #[serial]
    fn test_simple_single_quote() {
        env::set_var("KEY", "''");
        assert_eq!(&var("KEY").unwrap(), "''");
    }

    #[test]
    #[serial]
    fn test_single_quote_with_dollar() {
        env::set_var("KEY", "'$'");
        assert_eq!(&var("KEY").unwrap(), "'$'");
    }

    #[test]
    #[serial]
    fn test_single_quote_with_not_var() {
        env::set_var("KEY", "'${KEY2}'");
        env::set_var("KEY2", "bad");
        assert_eq!(&var("KEY").unwrap(), "'${KEY2}'");
    }

    #[test]
    #[serial]
    fn test_single_quote_with_not_var_no_braces() {
        env::set_var("KEY", "'$KEY2'");
        env::set_var("KEY2", "bad");
        assert_eq!(&var("KEY").unwrap(), "'$KEY2'");
    }

    #[test]
    #[serial]
    fn test_single_quote_with_non_matching_brace() {
        env::set_var("KEY", "'}'");
        assert_eq!(&var("KEY").unwrap(), "'}'");
    }

    #[test]
    #[serial]
    fn test_single_quotes_encapsulating_quote() {
        env::set_var("KEY", "'\"'");
        assert_eq!(&var("KEY").unwrap(), "'\"'");
    }

    #[test]
    #[serial]
    fn test_some_nonrecursive_wierd_ones_from_my_env() {
        let vars = vec![
            ("is_vim", "ps -o state= -o comm= -t '#{pane_tty}'     | grep -iqE '^[^TXZ ]+ +(\\S+\\/)?g?(view|n?vim?x?)(diff)?$'"),
            ("tmux_version", "$(tmux -V | sed -En \"s/^tmux ([0-9]+(.[0-9]+)?).*/\\1/p\")"),
        ];
        for (k, v) in vars {
            env::set_var(k, v);
            assert_eq!(&var(k).unwrap(), v);
        }
    }

    #[test]
    #[serial]
    fn test_trim_quotes() {
        let key1 = "KEY1";
        let key2 = "KEY2";
        env::set_var(key2, "test");
        env::set_var(key1, "\"${KEY2}\"");
        assert_eq!(var(key1).unwrap(), "test".to_string());
    }

    #[test]
    #[serial]
    fn test_trim_quotes_with_random_padding_chars() {
        let key1 = "KEY1";
        let key2 = "KEY2";
        env::set_var(key2, "test");
        env::set_var(key1, "--\"${KEY2}\"--");
        assert_eq!(var(key1).unwrap(), "--test--".to_string());
    }

    #[test]
    #[serial]
    fn test_recursive_and_trim_quotes() {
        let key1 = "KEY1";
        let key2 = "KEY2";
        env::set_var(key2, "number2");
        env::set_var(key1, "number1 \"${KEY2}\" number3");
        assert_eq!(var(key2).unwrap(), "number2".to_string());
        assert_eq!(var(key1).unwrap(), "number1 number2 number3".to_string());
    }

    #[test]
    #[serial]
    fn test_more_recursive_and_trim_quotes() {
        let key1 = "KEY1";
        let key2 = "KEY2";
        let key3 = "KEY3";
        env::set_var(key3, "number3");
        env::set_var(key2, "number2 \"${KEY3}\" number4");
        env::set_var(key1, "number1 \"${KEY2}\" number5");
        assert_eq!(var(key3).unwrap(), "number3".to_string());
        assert_eq!(var(key2).unwrap(), "number2 number3 number4".to_string());
        assert_eq!(
            var(key1).unwrap(),
            "number1 number2 number3 number4 number5".to_string()
        );
    }

    #[test]
    #[serial]
    fn test_stop_no_brace_var_on_quotes() {
        let key1 = "KEY1";
        let key2 = "KEY2";
        env::set_var(key2, "number2");
        env::set_var(key1, "number1\"$KEY2\"number3");
        assert_eq!(var(key2).unwrap(), "number2".to_string());
        assert_eq!(var(key1).unwrap(), "number1number2number3".to_string());
    }

    #[test]
    #[serial]
    fn test_stop_no_brace_var_on_quotes_more_recursive() {
        let key1 = "KEY1";
        let key2 = "KEY2";
        let key3 = "KEY3";
        env::set_var(key3, "number3");
        env::set_var(key2, "number2\"$KEY3\"number4");
        env::set_var(key1, "number1\"$KEY2\"number5");
        assert_eq!(var(key3).unwrap(), "number3".to_string());
        assert_eq!(var(key2).unwrap(), "number2number3number4".to_string());
        assert_eq!(
            var(key1).unwrap(),
            "number1number2number3number4number5".to_string()
        );
    }
}