improvement: add parser for format strings (#1021)

This PR implements the parser of format strings described in #624.

10 files changed, 383 insertions, 4 deletions

diff --git a/src/config.rs b/src/config.rs
index d3f868f16..235351556 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -305,7 +305,7 @@ impl Default for SegmentConfig<'static> {
  - 'italic'
  - '<color>'        (see the parse_color_string doc for valid color strings)
 */
-fn parse_style_string(style_string: &str) -> Option<ansi_term::Style> {
+pub fn parse_style_string(style_string: &str) -> Option<ansi_term::Style> {
     style_string
         .split_whitespace()
         .fold(Some(ansi_term::Style::new()), |maybe_style, token| {
diff --git a/src/formatter/mod.rs b/src/formatter/mod.rs
new file mode 100644
index 000000000..b0607df0a
--- /dev/null
+++ b/src/formatter/mod.rs
@@ -0,0 +1,5 @@
+pub mod model;
+mod parser;
+pub mod string_formatter;
+
+pub use string_formatter::StringFormatter;
diff --git a/src/formatter/model.rs b/src/formatter/model.rs
new file mode 100644
index 000000000..8e4bbea39
--- /dev/null
+++ b/src/formatter/model.rs
@@ -0,0 +1,17 @@
+use std::borrow::Cow;
+
+pub struct TextGroup<'a> {
+    pub format: Vec<FormatElement<'a>>,
+    pub style: Vec<StyleElement<'a>>,
+}
+
+pub enum FormatElement<'a> {
+    Text(Cow<'a, str>),
+    Variable(Cow<'a, str>),
+    TextGroup(TextGroup<'a>),
+}
+
+pub enum StyleElement<'a> {
+    Text(Cow<'a, str>),
+    Variable(Cow<'a, str>),
+}
diff --git a/src/formatter/parser.rs b/src/formatter/parser.rs
new file mode 100644
index 000000000..95c267a3a
--- /dev/null
+++ b/src/formatter/parser.rs
@@ -0,0 +1,76 @@
+use pest::{error::Error, iterators::Pair, Parser};
+
+use super::model::*;
+
+#[derive(Parser)]
+#[grammar = "formatter/spec.pest"]
+struct IdentParser;
+
+fn _parse_textgroup(textgroup: Pair<Rule>) -> TextGroup {
+    let mut inner_rules = textgroup.into_inner();
+    let format = inner_rules.next().unwrap();
+    let style = inner_rules.next().unwrap();
+
+    TextGroup {
+        format: _parse_format(format),
+        style: _parse_style(style),
+    }
+}
+
+fn _parse_variable(variable: Pair<Rule>) -> &str {
+    variable.into_inner().next().unwrap().as_str()
+}
+
+fn _parse_text(text: Pair<Rule>) -> String {
+    let mut result = String::new();
+    for pair in text.into_inner() {
+        result.push_str(pair.as_str());
+    }
+    result
+}
+
+fn _parse_format(format: Pair<Rule>) -> Vec<FormatElement> {
+    let mut result: Vec<FormatElement> = Vec::new();
+
+    for pair in format.into_inner() {
+        match pair.as_rule() {
+            Rule::text => result.push(FormatElement::Text(_parse_text(pair).into())),
+            Rule::variable => result.push(FormatElement::Variable(_parse_variable(pair).into())),
+            Rule::textgroup => result.push(FormatElement::TextGroup(_parse_textgroup(pair))),
+            _ => unreachable!(),
+        }
+    }
+
+    result
+}
+
+fn _parse_style(style: Pair<Rule>) -> Vec<StyleElement> {
+    let mut result: Vec<StyleElement> = Vec::new();
+
+    for pair in style.into_inner() {
+        match pair.as_rule() {
+            Rule::text => result.push(StyleElement::Text(_parse_text(pair).into())),
+            Rule::variable => result.push(StyleElement::Variable(_parse_variable(pair).into())),
+            _ => unreachable!(),
+        }
+    }
+
+    result
+}
+
+pub fn parse(format: &str) -> Result<Vec<FormatElement>, Error<Rule>> {
+    let pairs = IdentParser::parse(Rule::expression, format)?;
+    let mut result: Vec<FormatElement> = Vec::new();
+
+    // Lifetime of Segment is the same as result
+    for pair in pairs.take_while(|pair| pair.as_rule() != Rule::EOI) {
+        match pair.as_rule() {
+            Rule::text => result.push(FormatElement::Text(_parse_text(pair).into())),
+            Rule::variable => result.push(FormatElement::Variable(_parse_variable(pair).into())),
+            Rule::textgroup => result.push(FormatElement::TextGroup(_parse_textgroup(pair))),
+            _ => unreachable!(),
+        }
+    }
+
+    Ok(result)
+}
diff --git a/src/formatter/spec.pest b/src/formatter/spec.pest
new file mode 100644
index 000000000..36be53be5
--- /dev/null
+++ b/src/formatter/spec.pest
@@ -0,0 +1,16 @@
+expression = _{ SOI ~ value* ~ EOI }
+value = _{ text | variable | textgroup }
+
+variable = { "$" ~ variable_name }
+variable_name = @{ char+ }
+char = _{ 'a'..'z' | 'A'..'Z' | '0'..'9' | "_" }
+
+text = { text_inner+ }
+text_inner = _{ text_inner_char | escape }
+text_inner_char = { !("[" | "]" | "(" | ")" | "$" | "\\") ~ ANY }
+escape = _{ "\\" ~ escaped_char }
+escaped_char = { "[" | "]" | "(" | ")" | "\\" | "$" }
+
+textgroup = { "[" ~ format ~ "]" ~ "(" ~ style ~ ")" }
+format = { (variable | text | textgroup)* }
+style = { (variable | text)* }
diff --git a/src/formatter/string_formatter.rs b/src/formatter/string_formatter.rs
new file mode 100644
index 000000000..70938973f
--- /dev/null
+++ b/src/formatter/string_formatter.rs
@@ -0,0 +1,252 @@
+use ansi_term::Style;
+use pest::error::Error;
+use rayon::prelude::*;
+use std::collections::BTreeMap;
+
+use crate::config::parse_style_string;
+use crate::segment::Segment;
+
+use super::model::*;
+use super::parser::{parse, Rule};
+
+type VariableMapType = BTreeMap<String, Option<Vec<Segment>>>;
+
+pub struct StringFormatter<'a> {
+    format: Vec<FormatElement<'a>>,
+    variables: VariableMapType,
+}
+
+impl<'a> StringFormatter<'a> {
+    /// Creates an instance of StringFormatter from a format string
+    pub fn new(format: &'a str) -> Result<Self, Error<Rule>> {
+        parse(format)
+            .map(|format| {
+                let variables = _get_variables(&format);
+                (format, variables)
+            })
+            .map(|(format, variables)| Self { format, variables })
+    }
+
+    /// Maps variable name to its value
+    pub fn map(mut self, mapper: impl Fn(&str) -> Option<String> + Sync) -> Self {
+        self.variables.par_iter_mut().for_each(|(key, value)| {
+            *value = mapper(key).map(|value| vec![_new_segment(key.to_string(), value, None)]);
+        });
+        self
+    }
+
+    /// Maps variable name to an array of segments
+    pub fn map_variables_to_segments(
+        mut self,
+        mapper: impl Fn(&str) -> Option<Vec<Segment>> + Sync,
+    ) -> Self {
+        self.variables.par_iter_mut().for_each(|(key, value)| {
+            *value = mapper(key);
+        });
+        self
+    }
+
+    /// Parse the format string and consume self.
+    pub fn parse(self, default_style: Option<Style>) -> Vec<Segment> {
+        fn _parse_textgroup<'a>(
+            textgroup: TextGroup<'a>,
+            variables: &'a VariableMapType,
+        ) -> Vec<Segment> {
+            let style = _parse_style(textgroup.style);
+            _parse_format(textgroup.format, style, &variables)
+        }
+
+        fn _parse_style(style: Vec<StyleElement>) -> Option<Style> {
+            let style_string = style
+                .iter()
+                .flat_map(|style| match style {
+                    StyleElement::Text(text) => text.as_ref().chars(),
+                    StyleElement::Variable(variable) => {
+                        log::warn!(
+                            "Variable `{}` monitored in style string, which is not allowed",
+                            &variable
+                        );
+                        "".chars()
+                    }
+                })
+                .collect::<String>();
+            parse_style_string(&style_string)
+        }
+
+        fn _parse_format<'a>(
+            mut format: Vec<FormatElement<'a>>,
+            style: Option<Style>,
+            variables: &'a VariableMapType,
+        ) -> Vec<Segment> {
+            let mut result: Vec<Segment> = Vec::new();
+
+            format.reverse();
+            while let Some(el) = format.pop() {
+                let mut segments = match el {
+                    FormatElement::Text(text) => {
+                        vec![_new_segment("_text".into(), text.into_owned(), style)]
+                    }
+                    FormatElement::TextGroup(textgroup) => {
+                        let textgroup = TextGroup {
+                            format: textgroup.format,
+                            style: textgroup.style,
+                        };
+                        _parse_textgroup(textgroup, &variables)
+                    }
+                    FormatElement::Variable(name) => variables
+                        .get(name.as_ref())
+                        .map(|segments| segments.clone().unwrap_or_default())
+                        .unwrap_or_default(),
+                };
+                result.append(&mut segments);
+            }
+
+            result
+        }
+
+        _parse_format(self.format, default_style, &self.variables)
+    }
+}
+
+/// Extract variable names from an array of `FormatElement` into a `BTreeMap`
+fn _get_variables<'a>(format: &[FormatElement<'a>]) -> VariableMapType {
+    let mut variables: VariableMapType = Default::default();
+
+    fn _push_variables_from_textgroup<'a>(
+        variables: &mut VariableMapType,
+        textgroup: &'a TextGroup<'a>,
+    ) {
+        for el in &textgroup.format {
+            match el {
+                FormatElement::Variable(name) => _push_variable(variables, name.as_ref()),
+                FormatElement::TextGroup(textgroup) => {
+                    _push_variables_from_textgroup(variables, &textgroup)
+                }
+                _ => {}
+            }
+        }
+        for el in &textgroup.style {
+            if let StyleElement::Variable(name) = el {
+                _push_variable(variables, name.as_ref())
+            }
+        }
+    }
+
+    fn _push_variable<'a>(variables: &mut VariableMapType, name: &'a str) {
+        variables.insert(name.to_owned(), None);
+    }
+
+    for el in format {
+        match el {
+            FormatElement::Variable(name) => _push_variable(&mut variables, name.as_ref()),
+            FormatElement::TextGroup(textgroup) => {
+                _push_variables_from_textgroup(&mut variables, &textgroup)
+            }
+            _ => {}
+        }
+    }
+
+    variables
+}
+
+/// Helper function to create a new segment
+fn _new_segment(name: String, value: String, style: Option<Style>) -> Segment {
+    Segment {
+        _name: name,
+        value,
+        style,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use ansi_term::Color;
+
+    // match_next(result: Iter<Segment>, value, style)
+    macro_rules! match_next {
+        ($iter:ident, $value:literal, $($style:tt)+) => {
+            let _next = $iter.next().unwrap();
+            assert_eq!(_next.value, $value);
+            assert_eq!(_next.style, $($style)+);
+        }
+    }
+
+    fn empty_mapper(_: &str) -> Option<String> {
+        None
+    }
+
+    #[test]
+    fn test_default_style() {
+        const FORMAT_STR: &str = "text";
+        let style = Some(Color::Red.bold());
+
+        let formatter = StringFormatter::new(FORMAT_STR).unwrap().map(empty_mapper);
+        let result = formatter.parse(style);
+        let mut result_iter = result.iter();
+        match_next!(result_iter, "text", style);
+    }
+
+    #[test]
+    fn test_textgroup_text_only() {
+        const FORMAT_STR: &str = "[text](red bold)";
+        let formatter = StringFormatter::new(FORMAT_STR).unwrap().map(empty_mapper);
+        let result = formatter.parse(None);
+        let mut result_iter = result.iter();
+        match_next!(result_iter, "text", Some(Color::Red.bold()));
+    }
+
+    #[test]
+    fn test_variable_only() {
+        const FORMAT_STR: &str = "$var1";
+
+        let formatter = StringFormatter::new(FORMAT_STR)
+            .unwrap()
+            .map(|variable| match variable {
+                "var1" => Some("text1".to_owned()),
+                _ => None,
+            });
+        let result = formatter.parse(None);
+        let mut result_iter = result.iter();
+        match_next!(result_iter, "text1", None);
+    }
+
+    #[test]
+    fn test_escaped_chars() {
+        const FORMAT_STR: &str = r#"\\\[\$text\]\(red bold\)"#;
+
+        let formatter = StringFormatter::new(FORMAT_STR).unwrap().map(empty_mapper);
+        let result = formatter.parse(None);
+        let mut result_iter = result.iter();
+        match_next!(result_iter, r#"\[$text](red bold)"#, None);
+    }
+
+    #[test]
+    fn test_nested_textgroup() {
+        const FORMAT_STR: &str = "outer [middle [inner](blue)](red bold)";
+        let outer_style = Some(Color::Green.normal());
+        let middle_style = Some(Color::Red.bold());
+        let inner_style = Some(Color::Blue.normal());
+
+        let formatter = StringFormatter::new(FORMAT_STR).unwrap().map(empty_mapper);
+        let result = formatter.parse(outer_style);
+        let mut result_iter = result.iter();
+        match_next!(result_iter, "outer ", outer_style);
+        match_next!(result_iter, "middle ", middle_style);
+        match_next!(result_iter, "inner", inner_style);
+    }
+
+    #[test]
+    fn test_parse_error() {
+        // brackets without escape
+        {
+            const FORMAT_STR: &str = "[";
+            assert!(StringFormatter::new(FORMAT_STR).is_err());
+        }
+        // Dollar without variable
+        {
+            const FORMAT_STR: &str = "$ ";
+            assert!(StringFormatter::new(FORMAT_STR).is_err());
+        }
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index 08c91ec3a..7621ecd4e 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,7 +1,11 @@
+#[macro_use]
+extern crate pest_derive;
+
 // Lib is present to allow for benchmarking
 pub mod config;
 pub mod configs;
 pub mod context;
+pub mod formatter;
 pub mod module;
 pub mod modules;
 pub mod print;
diff --git a/src/main.rs b/src/main.rs
index 9b917a34b..1ae1eca0a 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,12 +2,15 @@ use std::time::SystemTime;
 
 #[macro_use]
 extern crate clap;
+#[macro_use]
+extern crate pest_derive;
 
 mod bug_report;
 mod config;
 mod configs;
 mod configure;
 mod context;
+mod formatter;
 mod init;
 mod module;
 mod modules;
diff --git a/src/module.rs b/src/module.rs
index 7fc2c72e5..79a4b6407 100644
--- a/src/module.rs
+++ b/src/module.rs
@@ -99,6 +99,11 @@ impl<'a> Module<'a> {
         self.segments.last_mut().unwrap()
     }
 
+    /// Set segments in module
+    pub fn set_segment(&mut self, segments: Vec<Segment>) {
+        self.segments = segments;
+    }
+
     /// Get module's name
     pub fn get_name(&self) -> &String {
         &self._name
diff --git a/src/segment.rs b/src/segment.rs
index d01034bb0..73e24b9a8 100644
--- a/src/segment.rs
+++ b/src/segment.rs
@@ -4,15 +4,16 @@ use std::fmt;
 /// A segment is a single configurable element in a module. This will usually
 /// contain a data point to provide context for the prompt's user
 /// (e.g. The version that software is running).
+#[derive(Clone)]
 pub struct Segment {
     /// The segment's name, to be used in configuration and logging.
-    _name: String,
+    pub _name: String,
 
     /// The segment's style. If None, will inherit the style of the module containing it.
-    style: Option<Style>,
+    pub style: Option<Style>,
 
     /// The string value of the current segment.
-    value: String,
+    pub value: String,
 }
 
 impl Segment {