feat: Add a fill module to pad out the line (#3029)

11 files changed, 304 insertions, 46 deletions

diff --git a/src/configs/fill.rs b/src/configs/fill.rs
new file mode 100644
index 000000000..f0bb468fa
--- /dev/null
+++ b/src/configs/fill.rs
@@ -0,0 +1,19 @@
+use crate::config::ModuleConfig;
+
+use serde::Serialize;
+use starship_module_config_derive::ModuleConfig;
+
+#[derive(Clone, ModuleConfig, Serialize)]
+pub struct FillConfig<'a> {
+    pub style: &'a str,
+    pub symbol: &'a str,
+}
+
+impl<'a> Default for FillConfig<'a> {
+    fn default() -> Self {
+        FillConfig {
+            style: "bold black",
+            symbol: ".",
+        }
+    }
+}
diff --git a/src/configs/mod.rs b/src/configs/mod.rs
index f403da008..5e844fb36 100644
--- a/src/configs/mod.rs
+++ b/src/configs/mod.rs
@@ -21,6 +21,7 @@ pub mod elixir;
 pub mod elm;
 pub mod env_var;
 pub mod erlang;
+pub mod fill;
 pub mod gcloud;
 pub mod git_branch;
 pub mod git_commit;
@@ -96,6 +97,7 @@ pub struct FullConfig<'a> {
     elm: elm::ElmConfig<'a>,
     env_var: IndexMap<String, env_var::EnvVarConfig<'a>>,
     erlang: erlang::ErlangConfig<'a>,
+    fill: fill::FillConfig<'a>,
     gcloud: gcloud::GcloudConfig<'a>,
     git_branch: git_branch::GitBranchConfig<'a>,
     git_commit: git_commit::GitCommitConfig<'a>,
@@ -169,6 +171,7 @@ impl<'a> Default for FullConfig<'a> {
             elm: Default::default(),
             env_var: Default::default(),
             erlang: Default::default(),
+            fill: Default::default(),
             gcloud: Default::default(),
             git_branch: Default::default(),
             git_commit: Default::default(),
diff --git a/src/context.rs b/src/context.rs
index 680d18cab..9f460dc51 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -49,6 +49,9 @@ pub struct Context<'a> {
     /// Construct the right prompt instead of the left prompt
     pub right: bool,
 
+    /// Width of terminal, or zero if width cannot be detected.
+    pub width: usize,
+
     /// A HashMap of environment variable mocks
     #[cfg(test)]
     pub env: HashMap<&'a str, String>,
@@ -135,6 +138,9 @@ impl<'a> Context<'a> {
             repo: OnceCell::new(),
             shell,
             right,
+            width: term_size::dimensions()
+                .map(|(width, _)| width)
+                .unwrap_or_default(),
             #[cfg(test)]
             env: HashMap::new(),
             #[cfg(test)]
diff --git a/src/formatter/string_formatter.rs b/src/formatter/string_formatter.rs
index 1a570ee9f..7778bb6ad 100644
--- a/src/formatter/string_formatter.rs
+++ b/src/formatter/string_formatter.rs
@@ -257,7 +257,7 @@ impl<'a> StringFormatter<'a> {
                 .into_iter()
                 .map(|el| {
                     match el {
-                        FormatElement::Text(text) => Ok(vec![Segment::new(style, text)]),
+                        FormatElement::Text(text) => Ok(Segment::from_text(style, text)),
                         FormatElement::TextGroup(textgroup) => {
                             let textgroup = TextGroup {
                                 format: textgroup.format,
@@ -274,13 +274,11 @@ impl<'a> StringFormatter<'a> {
                                     .into_iter()
                                     .map(|mut segment| {
                                         // Derive upper style if the style of segments are none.
-                                        if segment.style.is_none() {
-                                            segment.style = style;
-                                        };
+                                        segment.set_style_if_empty(style);
                                         segment
                                     })
                                     .collect()),
-                                VariableValue::Plain(text) => Ok(vec![Segment::new(style, text)]),
+                                VariableValue::Plain(text) => Ok(Segment::from_text(style, text)),
                                 VariableValue::Meta(format) => {
                                     let formatter = StringFormatter {
                                         format,
@@ -322,9 +320,9 @@ impl<'a> StringFormatter<'a> {
                                                     VariableValue::Plain(plain_value) => {
                                                         !plain_value.is_empty()
                                                     }
-                                                    VariableValue::Styled(segments) => {
-                                                        segments.iter().any(|x| !x.value.is_empty())
-                                                    }
+                                                    VariableValue::Styled(segments) => segments
+                                                        .iter()
+                                                        .any(|x| !x.value().is_empty()),
                                                 })
                                         })
                                 })
@@ -391,8 +389,8 @@ mod tests {
     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)+);
+            assert_eq!(_next.value(), $value);
+            assert_eq!(_next.style(), $($style)+);
         }
     }
 
@@ -511,14 +509,18 @@ mod tests {
         let styled_style = Some(Color::Green.italic());
         let styled_no_modifier_style = Some(Color::Green.normal());
 
+        let mut segments: Vec<Segment> = Vec::new();
+        segments.extend(Segment::from_text(None, "styless"));
+        segments.extend(Segment::from_text(styled_style, "styled"));
+        segments.extend(Segment::from_text(
+            styled_no_modifier_style,
+            "styled_no_modifier",
+        ));
+
         let formatter = StringFormatter::new(FORMAT_STR)
             .unwrap()
             .map_variables_to_segments(|variable| match variable {
-                "var" => Some(Ok(vec![
-                    Segment::new(None, "styless"),
-                    Segment::new(styled_style, "styled"),
-                    Segment::new(styled_no_modifier_style, "styled_no_modifier"),
-                ])),
+                "var" => Some(Ok(segments.clone())),
                 _ => None,
             });
         let result = formatter.parse(None).unwrap();
diff --git a/src/formatter/version.rs b/src/formatter/version.rs
index 808fa15ff..2f26d22c0 100644
--- a/src/formatter/version.rs
+++ b/src/formatter/version.rs
@@ -56,7 +56,7 @@ impl<'a> VersionFormatter<'a> {
         formatted.map(|segments| {
             segments
                 .iter()
-                .map(|segment| segment.value.as_str())
+                .map(|segment| segment.value())
                 .collect::<String>()
         })
     }
diff --git a/src/module.rs b/src/module.rs
index b81109320..2c88b5533 100644
--- a/src/module.rs
+++ b/src/module.rs
@@ -1,5 +1,5 @@
 use crate::context::Shell;
-use crate::segment::Segment;
+use crate::segment::{FillSegment, Segment};
 use crate::utils::wrap_colorseq_for_shell;
 use ansi_term::{ANSIString, ANSIStrings};
 use std::fmt;
@@ -26,6 +26,7 @@ pub const ALL_MODULES: &[&str] = &[
     "elm",
     "env_var",
     "erlang",
+    "fill",
     "gcloud",
     "git_branch",
     "git_commit",
@@ -124,29 +125,29 @@ impl<'a> Module<'a> {
         self.segments
             .iter()
             // no trim: if we add spaces/linebreaks it's not "empty" as we change the final output
-            .all(|segment| segment.value.is_empty())
+            .all(|segment| segment.value().is_empty())
     }
 
     /// Get values of the module's segments
     pub fn get_segments(&self) -> Vec<&str> {
         self.segments
             .iter()
-            .map(|segment| segment.value.as_str())
+            .map(|segment| segment.value())
             .collect()
     }
 
     /// Returns a vector of colored ANSIString elements to be later used with
     /// `ANSIStrings()` to optimize ANSI codes
     pub fn ansi_strings(&self) -> Vec<ANSIString> {
-        self.ansi_strings_for_shell(Shell::Unknown)
+        self.ansi_strings_for_shell(Shell::Unknown, None)
     }
 
-    pub fn ansi_strings_for_shell(&self, shell: Shell) -> Vec<ANSIString> {
-        let ansi_strings = self
-            .segments
-            .iter()
-            .map(Segment::ansi_string)
-            .collect::<Vec<ANSIString>>();
+    pub fn ansi_strings_for_shell(&self, shell: Shell, width: Option<usize>) -> Vec<ANSIString> {
+        let mut iter = self.segments.iter().peekable();
+        let mut ansi_strings: Vec<ANSIString> = Vec::new();
+        while iter.peek().is_some() {
+            ansi_strings.extend(ansi_line(&mut iter, width));
+        }
 
         match shell {
             Shell::Bash => ansi_strings_modified(ansi_strings, shell),
@@ -174,6 +175,49 @@ fn ansi_strings_modified(ansi_strings: Vec<ANSIString>, shell: Shell) -> Vec<ANS
         .collect::<Vec<ANSIString>>()
 }
 
+fn ansi_line<'a, I>(segments: &mut I, term_width: Option<usize>) -> Vec<ANSIString<'a>>
+where
+    I: Iterator<Item = &'a Segment>,
+{
+    let mut used = 0usize;
+    let mut current: Vec<ANSIString> = Vec::new();
+    let mut chunks: Vec<(Vec<ANSIString>, &FillSegment)> = Vec::new();
+
+    for segment in segments {
+        match segment {
+            Segment::Fill(fs) => {
+                chunks.push((current, fs));
+                current = Vec::new();
+            }
+            _ => {
+                used += segment.width_graphemes();
+                current.push(segment.ansi_string());
+            }
+        }
+
+        if let Segment::LineTerm = segment {
+            break;
+        }
+    }
+
+    if chunks.is_empty() {
+        current
+    } else {
+        let fill_size = term_width
+            .map(|tw| if tw > used { Some(tw - used) } else { None })
+            .flatten()
+            .map(|remaining| remaining / chunks.len());
+        chunks
+            .into_iter()
+            .flat_map(|(strs, fill)| {
+                strs.into_iter()
+                    .chain(std::iter::once(fill.ansi_string(fill_size)))
+            })
+            .chain(current.into_iter())
+            .collect::<Vec<ANSIString>>()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -208,7 +252,7 @@ mod tests {
             config: None,
             name: name.to_string(),
             description: desc.to_string(),
-            segments: vec![Segment::new(None, "")],
+            segments: Segment::from_text(None, ""),
             duration: Duration::default(),
         };
 
@@ -223,7 +267,7 @@ mod tests {
             config: None,
             name: name.to_string(),
             description: desc.to_string(),
-            segments: vec![Segment::new(None, "\n")],
+            segments: Segment::from_text(None, "\n"),
             duration: Duration::default(),
         };
 
@@ -238,7 +282,7 @@ mod tests {
             config: None,
             name: name.to_string(),
             description: desc.to_string(),
-            segments: vec![Segment::new(None, " ")],
+            segments: Segment::from_text(None, " "),
             duration: Duration::default(),
         };
 
diff --git a/src/modules/fill.rs b/src/modules/fill.rs
new file mode 100644
index 000000000..edbabbee9
--- /dev/null
+++ b/src/modules/fill.rs
@@ -0,0 +1,38 @@
+use super::{Context, Module};
+
+use crate::config::{parse_style_string, RootModuleConfig};
+use crate::configs::fill::FillConfig;
+use crate::segment::Segment;
+
+/// Creates a module that fills the any extra space on the line.
+///
+pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
+    let mut module = context.new_module("fill");
+    let config: FillConfig = FillConfig::try_load(module.config);
+
+    let style = parse_style_string(config.style);
+
+    module.set_segments(vec![Segment::fill(style, config.symbol)]);
+
+    Some(module)
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::test::ModuleRenderer;
+    use ansi_term::Color;
+
+    #[test]
+    fn basic() {
+        let actual = ModuleRenderer::new("fill")
+            .config(toml::toml! {
+                [fill]
+                style = "bold green"
+                symbol = "*-"
+            })
+            .collect();
+        let expected = Some(format!("{}", Color::Green.bold().paint("*-")));
+
+        assert_eq!(expected, actual);
+    }
+}
diff --git a/src/modules/line_break.rs b/src/modules/line_break.rs
index 4aa67595f..233e5f468 100644
--- a/src/modules/line_break.rs
+++ b/src/modules/line_break.rs
@@ -3,11 +3,9 @@ use crate::segment::Segment;
 
 /// Creates a module for the line break
 pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
-    const LINE_ENDING: &str = "\n";
-
     let mut module = context.new_module("line_break");
 
-    module.set_segments(vec![Segment::new(None, LINE_ENDING)]);
+    module.set_segments(vec![Segment::LineTerm]);
 
     Some(module)
 }
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index ccc360705..e3897aeb3 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -16,6 +16,7 @@ mod elixir;
 mod elm;
 mod env_var;
 mod erlang;
+mod fill;
 mod gcloud;
 mod git_branch;
 mod git_commit;
@@ -97,6 +98,7 @@ pub fn handle<'a>(module: &str, context: &'a Context) -> Option<Module<'a>> {
             "elm" => elm::module(context),
             "erlang" => erlang::module(context),
             "env_var" => env_var::module(context),
+            "fill" => fill::module(context),
             "gcloud" => gcloud::module(context),
             "git_branch" => git_branch::module(context),
             "git_commit" => git_commit::module(context),
@@ -181,6 +183,7 @@ pub fn description(module: &str) -> &'static str {
         "dotnet" => "The relevant version of the .NET Core SDK for the current directory",
         "env_var" => "Displays the current value of a selected environment variable",
         "erlang" => "Current OTP version",
+        "fill" => "Fills the remaining space on the line with a pad string",
         "gcloud" => "The current GCP client configuration",
         "git_branch" => "The active branch of the repo in your current directory",
         "git_commit" => "The active commit (and tag if any) of the repo in your current directory",
diff --git a/src/print.rs b/src/print.rs
index a202cb86e..cc30f7c3c 100644
--- a/src/print.rs
+++ b/src/print.rs
@@ -16,7 +16,7 @@ use crate::module::ALL_MODULES;
 use crate::modules;
 use crate::segment::Segment;
 
-pub struct Grapheme<'a>(&'a str);
+pub struct Grapheme<'a>(pub &'a str);
 
 impl<'a> Grapheme<'a> {
     pub fn width(&self) -> usize {
@@ -112,7 +112,7 @@ pub fn get_prompt(context: Context) -> String {
             .expect("Unexpected error returned in root format variables"),
     );
 
-    let module_strings = root_module.ansi_strings_for_shell(context.shell);
+    let module_strings = root_module.ansi_strings_for_shell(context.shell, Some(context.width));
     if config.add_newline {
         writeln!(buf).unwrap();
     }
diff --git a/src/segment.rs b/src/segment.rs
index 200a93b9e..643298c2d 100644
--- a/src/segment.rs
+++ b/src/segment.rs
@@ -1,39 +1,184 @@
+use crate::print::{Grapheme, UnicodeWidthGraphemes};
 use ansi_term::{ANSIString, Style};
 use std::fmt;
+use unicode_segmentation::UnicodeSegmentation;
 
-/// 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).
+/// Type that holds text with an associated style
 #[derive(Clone)]
-pub struct Segment {
+pub struct TextSegment {
     /// The segment's style. If None, will inherit the style of the module containing it.
-    pub style: Option<Style>,
+    style: Option<Style>,
 
     /// The string value of the current segment.
-    pub value: String,
+    value: String,
+}
+
+impl TextSegment {
+    // Returns the ANSIString of the segment value
+    fn ansi_string(&self) -> ANSIString {
+        match self.style {
+            Some(style) => style.paint(&self.value),
+            None => ANSIString::from(&self.value),
+        }
+    }
+}
+
+/// Type that holds fill text with an associated style
+#[derive(Clone)]
+pub struct FillSegment {
+    /// The segment's style. If None, will inherit the style of the module containing it.
+    style: Option<Style>,
+
+    /// The string value of the current segment.
+    value: String,
+}
+
+impl FillSegment {
+    // Returns the ANSIString of the segment value, not including its prefix and suffix
+    pub fn ansi_string(&self, width: Option<usize>) -> ANSIString {
+        let s = match width {
+            Some(w) => self
+                .value
+                .graphemes(true)
+                .cycle()
+                .scan(0usize, |len, g| {
+                    *len += Grapheme(g).width();
+                    if *len <= w {
+                        Some(g)
+                    } else {
+                        None
+                    }
+                })
+                .collect::<String>(),
+            None => String::from(&self.value),
+        };
+        match self.style {
+            Some(style) => style.paint(s),
+            None => ANSIString::from(s),
+        }
+    }
+}
+
+#[cfg(test)]
+mod fill_seg_tests {
+    use super::FillSegment;
+    use ansi_term::Color;
+
+    #[test]
+    fn ansi_string_width() {
+        let width: usize = 10;
+        let style = Color::Blue.bold();
+
+        let inputs = vec![
+            (".", ".........."),
+            (".:", ".:.:.:.:.:"),
+            ("-:-", "-:--:--:--"),
+            ("🟦", "🟦🟦🟦🟦🟦"),
+            ("🟢🔵🟡", "🟢🔵🟡🟢🔵"),
+        ];
+
+        for (text, expected) in inputs.iter() {
+            let f = FillSegment {
+                value: String::from(*text),
+                style: Some(style),
+            };
+            let actual = f.ansi_string(Some(width));
+            assert_eq!(style.paint(*expected), actual);
+        }
+    }
+}
+
+/// A segment is a styled text chunk ready for printing.
+#[derive(Clone)]
+pub enum Segment {
+    Text(TextSegment),
+    Fill(FillSegment),
+    LineTerm,
 }
 
 impl Segment {
-    /// Creates a new segment.
-    pub fn new<T>(style: Option<Style>, value: T) -> Self
+    /// Creates new segments from a text with a style; breaking out LineTerminators.
+    pub fn from_text<T>(style: Option<Style>, value: T) -> Vec<Segment>
+    where
+        T: Into<String>,
+    {
+        let mut segs: Vec<Segment> = Vec::new();
+        value.into().split(LINE_TERMINATOR).for_each(|s| {
+            if !segs.is_empty() {
+                segs.push(Segment::LineTerm)
+            }
+            segs.push(Segment::Text(TextSegment {
+                value: String::from(s),
+                style,
+            }))
+        });
+        segs
+    }
+
+    /// Creates a new fill segment
+    pub fn fill<T>(style: Option<Style>, value: T) -> Self
     where
         T: Into<String>,
     {
-        Self {
+        Segment::Fill(FillSegment {
             style,
             value: value.into(),
+        })
+    }
+
+    pub fn style(&self) -> Option<Style> {
+        match self {
+            Segment::Fill(fs) => fs.style,
+            Segment::Text(ts) => ts.style,
+            Segment::LineTerm => None,
+        }
+    }
+
+    pub fn set_style_if_empty(&mut self, style: Option<Style>) {
+        match self {
+            Segment::Fill(fs) => {
+                if fs.style.is_none() {
+                    fs.style = style
+                }
+            }
+            Segment::Text(ts) => {
+                if ts.style.is_none() {
+                    ts.style = style
+                }
+            }
+            Segment::LineTerm => {}
+        }
+    }
+
+    pub fn value(&self) -> &str {
+        match self {
+            Segment::Fill(fs) => &fs.value,
+            Segment::Text(ts) => &ts.value,
+            Segment::LineTerm => LINE_TERMINATOR_STRING,
         }
     }
 
     // Returns the ANSIString of the segment value, not including its prefix and suffix
     pub fn ansi_string(&self) -> ANSIString {
-        match self.style {
-            Some(style) => style.paint(&self.value),
-            None => ANSIString::from(&self.value),
+        match self {
+            Segment::Fill(fs) => fs.ansi_string(None),
+            Segment::Text(ts) => ts.ansi_string(),
+            Segment::LineTerm => ANSIString::from(LINE_TERMINATOR_STRING),
+        }
+    }
+
+    pub fn width_graphemes(&self) -> usize {
+        match self {
+            Segment::Fill(fs) => fs.value.width_graphemes(),
+            Segment::Text(ts) => ts.value.width_graphemes(),
+            Segment::LineTerm => 0,
         }
     }
 }
 
+const LINE_TERMINATOR: char = '\n';
+const LINE_TERMINATOR_STRING: &str = "\n";
+
 impl fmt::Display for Segment {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!(f, "{}", self.ansi_string())