use crate::parser::{file, FileLine}; use nom_locate::LocatedSpan; use std::env; use std::env::VarError; use std::fmt::Display; use std::str::FromStr; use std::string::ParseError; #[derive(Clone, Debug, Default)] pub struct EnvironmentFile { lines: Vec<FileLine>, } #[derive(Debug, thiserror::Error)] pub enum EnvironmentFileError { #[error("Unable to determine eof marker")] InvalidEof, #[error("Parse error at line {line}, column {column}: {kind:?}")] ParseError { line: usize, column: usize, kind: nom::error::ErrorKind, }, #[error(transparent)] StdParseError(#[from] ParseError), #[error("Value could not be determined to be valid UTF-8")] InvalidValue, } impl FromStr for EnvironmentFile { type Err = EnvironmentFileError; fn from_str(s: &str) -> Result<Self, Self::Err> { let (_, lines) = file(LocatedSpan::new(s)).map_err(|err| match err { nom::Err::Incomplete(_) => EnvironmentFileError::InvalidEof, nom::Err::Error(err) | nom::Err::Failure(err) => EnvironmentFileError::ParseError { column: err.input.get_column(), line: err.input.location_line() as usize, kind: err.code, }, })?; Ok(Self::new(lines)) } } fn is_kv_line(line: &FileLine) -> bool { matches!(line, FileLine::KeyValue { .. }) } fn is_comment_line(line: &FileLine) -> bool { matches!(line, FileLine::Comment(_)) } impl EnvironmentFile { pub fn new(lines: Vec<FileLine>) -> Self { Self { lines } } pub fn parse(value: impl AsRef<str>) -> Result<Self, EnvironmentFileError> { value.as_ref().parse() } pub fn iter(&self) -> EnvFileIterator { EnvFileIterator { lines: &self.lines, current: 0, } } pub fn lines_kv(&self) -> SubTypeFileIterator { SubTypeFileIterator { lines: &self.lines, current: 0, predicate: is_kv_line, } } pub fn lines_comment(&self) -> SubTypeFileIterator { SubTypeFileIterator { lines: &self.lines, current: 0, predicate: is_comment_line, } } pub fn get_raw(&self, key: &str) -> Option<String> { for line in &self.lines { if let FileLine::KeyValue { key: k, value } = line { if k == key { return Some(value.iter().map(|part| part.to_string()).collect::<String>()); } } } None } pub fn len(&self) -> usize { self.lines.len() } pub fn is_empty(&self) -> bool { self.lines.is_empty() } } #[derive(Clone, Debug, Default)] pub struct ApplyOptions { pub prefix: Option<String>, pub overwrite: bool, } impl ApplyOptions { pub fn new(prefix: impl Display, overwrite: bool) -> Self { Self { prefix: Some(prefix.to_string()), overwrite, } } pub fn with_prefix(prefix: impl Display) -> Self { Self::new(prefix, false) } pub fn with_overwrite(overwrite: bool) -> Self { Self { prefix: None, overwrite } } } pub trait ApplyEnvironmentFile { fn apply(&self, options: ApplyOptions); } impl ApplyEnvironmentFile for EnvironmentFile { fn apply(&self, options: ApplyOptions) { let _ = set_from_file(self, options); } } impl<E> ApplyEnvironmentFile for Result<EnvironmentFile, E> { fn apply(&self, options: ApplyOptions) { if let Ok(file) = self { file.apply(options); } } } impl ApplyEnvironmentFile for Option<EnvironmentFile> { fn apply(&self, options: ApplyOptions) { if let Some(file) = self { file.apply(options); } } } pub struct EnvFileIterator<'a> { lines: &'a [FileLine], current: usize, } impl<'a> Iterator for EnvFileIterator<'a> { type Item = &'a FileLine; fn next(&mut self) -> Option<Self::Item> { if self.current >= self.lines.len() { None } else { let line = &self.lines[self.current]; self.current += 1; Some(line) } } } impl ExactSizeIterator for EnvFileIterator<'_> { fn len(&self) -> usize { self.lines.len() - self.current } } pub struct SubTypeFileIterator<'a> { lines: &'a [FileLine], current: usize, predicate: fn(&FileLine) -> bool, } impl<'a> Iterator for SubTypeFileIterator<'a> { type Item = &'a FileLine; fn next(&mut self) -> Option<Self::Item> { let mut found = None; while self.current < self.lines.len() { let line = &self.lines[self.current]; self.current += 1; if (self.predicate)(line) { found = Some(line); break; } } found } } fn is_missing_key(key: &str) -> Result<bool, EnvironmentFileError> { match env::var(key) { Ok(_) => Ok(false), Err(VarError::NotPresent) => Ok(true), Err(_) => Err(EnvironmentFileError::InvalidValue), } } fn set_from_file(file: &EnvironmentFile, options: ApplyOptions) -> Result<(), EnvironmentFileError> { let mut defferred = Vec::with_capacity(file.len()); for line in file.lines_kv() { if let FileLine::KeyValue { key, .. } = &line { let key_name = if let Some(prefix) = &options.prefix { format!("{}{}", prefix, key) } else { key.to_string() }; if options.overwrite || is_missing_key(&key_name)? { if line.is_complete() { eprintln!("Setting {} with value {}", key_name, line.assemble_value()); env::set_var(key_name, line.assemble_value()); } else { defferred.push(line); } } } } for line in defferred { if let FileLine::KeyValue { key, .. } = line { let key_name = if let Some(prefix) = &options.prefix { format!("{}{}", prefix, key) } else { key.to_string() }; env::set_var(key_name, line.assemble_value()); } } Ok(()) } #[cfg(test)] mod tests { use super::*; use crate::parser::ValuePart; #[test] fn test_lines_kv_iterator() { let env_file = EnvironmentFile::new(vec![ FileLine::empty(), FileLine::comment("# This is a comment"), FileLine::key_value("KEY1".to_string(), vec![ValuePart::Static("VALUE1".to_string())]), FileLine::empty(), FileLine::key_value("KEY2".to_string(), vec![ValuePart::Static("VALUE2".to_string())]), ]); let kv_lines: Vec<&FileLine> = env_file.lines_kv().collect(); assert_eq!(kv_lines.len(), 2); assert!( matches!(kv_lines[0], FileLine::KeyValue { key, value } if key == "KEY1" && value == &vec![ValuePart::Static("VALUE1".to_string())]) ); assert!( matches!(kv_lines[1], FileLine::KeyValue { key, value } if key == "KEY2" && value == &vec![ValuePart::Static("VALUE2".to_string())]) ); } #[test] fn test_lines_comment_iterator() { let env_file = EnvironmentFile::new(vec![ FileLine::empty(), FileLine::comment(" Comment 1"), FileLine::key_value("KEY1".to_string(), vec![ValuePart::Static("VALUE1".to_string())]), FileLine::comment(" Comment 2"), FileLine::key_value("KEY2".to_string(), vec![ValuePart::Static("VALUE2".to_string())]), FileLine::comment(" Comment 3"), ]); let comment_lines: Vec<&FileLine> = env_file.lines_comment().collect(); assert_eq!(comment_lines.len(), 3); assert!(matches!(comment_lines[0], FileLine::Comment(comment) if comment == " Comment 1")); assert!(matches!(comment_lines[1], FileLine::Comment(comment) if comment == " Comment 2")); assert!(matches!(comment_lines[2], FileLine::Comment(comment) if comment == " Comment 3")); } #[test] fn test_iter_returns_all_lines() { let env_file = EnvironmentFile::new(vec![ FileLine::empty(), FileLine::comment("This is a comment"), FileLine::key_value("KEY1".to_string(), vec![ValuePart::Static("VALUE1".to_string())]), FileLine::empty(), FileLine::key_value("KEY2".to_string(), vec![ValuePart::Static("VALUE2".to_string())]), ]); let all_lines: Vec<&FileLine> = env_file.iter().collect(); assert_eq!(all_lines.len(), 5); assert!(matches!(all_lines[0], FileLine::Empty)); assert!(matches!(all_lines[1], FileLine::Comment(comment) if comment == "This is a comment")); assert!( matches!(all_lines[2], FileLine::KeyValue { key, value } if key == "KEY1" && value == &vec![ValuePart::Static("VALUE1".to_string())]) ); assert!(matches!(all_lines[3], FileLine::Empty)); assert!( matches!(all_lines[4], FileLine::KeyValue { key, value } if key == "KEY2" && value == &vec![ValuePart::Static("VALUE2".to_string())]) ); } #[test] fn test_lines_kv_maintains_order() { let env_file = EnvironmentFile::new(vec![ FileLine::empty(), FileLine::key_value("KEY1".to_string(), vec![ValuePart::Static("VALUE1".to_string())]), FileLine::comment("Comment"), FileLine::key_value("KEY2".to_string(), vec![ValuePart::Static("VALUE2".to_string())]), FileLine::key_value("KEY3".to_string(), vec![ValuePart::Static("VALUE3".to_string())]), FileLine::empty(), ]); let kv_lines: Vec<&FileLine> = env_file.lines_kv().collect(); assert_eq!(kv_lines.len(), 3); assert!( matches!(kv_lines[0], FileLine::KeyValue { key, value } if key == "KEY1" && value == &vec![ValuePart::Static("VALUE1".to_string())]) ); assert!( matches!(kv_lines[1], FileLine::KeyValue { key, value } if key == "KEY2" && value == &vec![ValuePart::Static("VALUE2".to_string())]) ); assert!( matches!(kv_lines[2], FileLine::KeyValue { key, value } if key == "KEY3" && value == &vec![ValuePart::Static("VALUE3".to_string())]) ); } }