use ansi_term::Style; use pest::error::Error as PestError; use rayon::prelude::*; use std::borrow::Cow; use std::collections::{BTreeMap, BTreeSet}; use std::error::Error; use std::fmt; use crate::config::parse_style_string; use crate::context::{Context, Shell}; use crate::segment::Segment; use super::model::*; use super::parser::{parse, Rule}; #[derive(Clone)] enum VariableValue<'a> { Plain(Cow<'a, str>), NoEscapingPlain(Cow<'a, str>), Styled(Vec), Meta(Vec>), } impl<'a> Default for VariableValue<'a> { fn default() -> Self { VariableValue::Plain(Cow::Borrowed("")) } } type VariableMapType<'a> = BTreeMap, StringFormatterError>>>; type StyleVariableMapType<'a> = BTreeMap, StringFormatterError>>>; #[derive(Debug, Clone, PartialEq)] pub enum StringFormatterError { Custom(String), Parse(PestError), } impl fmt::Display for StringFormatterError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::Custom(error) => write!(f, "{}", error), Self::Parse(error) => write!(f, "{}", error), } } } impl Error for StringFormatterError {} impl From for StringFormatterError { fn from(error: String) -> Self { Self::Custom(error) } } pub struct StringFormatter<'a> { format: Vec>, variables: VariableMapType<'a>, style_variables: StyleVariableMapType<'a>, } impl<'a> StringFormatter<'a> { /// Creates an instance of StringFormatter from a format string /// /// This method will throw an Error when the given format string fails to parse. pub fn new(format: &'a str) -> Result { let format = parse(format).map_err(StringFormatterError::Parse)?; // Cache all variables let variables = format .get_variables() .into_iter() .map(|key| (key.to_string(), None)) .collect::(); let style_variables = format .get_style_variables() .into_iter() .map(|key| (key.to_string(), None)) .collect::(); Ok(Self { format, variables, style_variables, }) } /// A StringFormatter that does no formatting, parse just returns the raw text pub fn raw(text: &'a str) -> Self { Self { format: vec![FormatElement::Text(text.into())], variables: BTreeMap::new(), style_variables: BTreeMap::new(), } } /// Maps variable name to its value /// /// You should provide a function or closure that accepts the variable name `name: &str` as a /// parameter and returns the one of the following values: /// /// - `None`: This variable will be reserved for further mappers. If it is `None` when /// `self.parse()` is called, it will be dropped. /// /// - `Some(Err(StringFormatterError))`: This variable will throws `StringFormatterError` when /// `self.parse()` is called. Return this if some fatal error occurred and the format string /// should not be rendered. /// /// - `Some(Ok(_))`: The value of this variable will be displayed in the format string. /// pub fn map(mut self, mapper: M) -> Self where T: Into>, M: Fn(&str) -> Option> + Sync, { self.variables .par_iter_mut() .filter(|(_, value)| value.is_none()) .for_each(|(key, value)| { *value = mapper(key).map(|var| var.map(|var| VariableValue::Plain(var.into()))); }); self } /// Maps variable name into a value which is wrapped to prevent escaping later /// /// This should be used for variables that should not be escaped before inclusion in the prompt /// /// See `StringFormatter::map` for description on the parameters. /// pub fn map_no_escaping(mut self, mapper: M) -> Self where T: Into>, M: Fn(&str) -> Option> + Sync, { self.variables .par_iter_mut() .filter(|(_, value)| value.is_none()) .for_each(|(key, value)| { *value = mapper(key) .map(|var| var.map(|var| VariableValue::NoEscapingPlain(var.into()))); }); self } /// Maps a meta-variable to a format string containing other variables. /// /// This function should be called **before** other map methods so that variables found in /// the format strings of meta-variables can be cached properly. /// /// See `StringFormatter::map` for description on the parameters. pub fn map_meta(mut self, mapper: M) -> Self where M: Fn(&str, &BTreeSet) -> Option<&'a str> + Sync, { let variables = self.get_variables(); let (variables, style_variables) = self .variables .iter_mut() .filter(|(_, value)| value.is_none()) .fold( (VariableMapType::new(), StyleVariableMapType::new()), |(mut v, mut sv), (key, value)| { *value = mapper(key, &variables).map(|format| { StringFormatter::new(format).map(|formatter| { let StringFormatter { format, mut variables, mut style_variables, } = formatter; // Add variables in meta variables to self v.append(&mut variables); sv.append(&mut style_variables); VariableValue::Meta(format) }) }); (v, sv) }, ); self.variables.extend(variables); self.style_variables.extend(style_variables); self } /// Maps variable name to an array of segments /// /// See `StringFormatter::map` for description on the parameters. pub fn map_variables_to_segments(mut self, mapper: M) -> Self where M: Fn(&str) -> Option, StringFormatterError>> + Sync, { self.variables .par_iter_mut() .filter(|(_, value)| value.is_none()) .for_each(|(key, value)| { *value = mapper(key).map(|var| var.map(VariableValue::Styled)); }); self } /// Maps variable name in a style string to its value /// /// See `StringFormatter::map` for description on the parameters. pub fn map_style(mut self, mapper: M) -> Self where T: Into>, M: Fn(&str) -> Option> + Sync, { self.style_variables .par_iter_mut() .filter(|(_, value)| value.is_none()) .for_each(|(key, value)| { *value = mapper(key).map(|var| var.map(std::convert::Into::into)); }); self } /// Parse the format string and consume self. /// /// This method will throw an Error in the following conditions: /// /// - Format string in meta variables fails to parse /// - Variable mapper returns an error. pub fn parse( self, default_style: Option