Make segments optionals

7 files changed, 36 insertions, 22 deletions

diff --git a/src/modules/character.rs b/src/modules/character.rs
index 681c1af02..53a43e886 100644
--- a/src/modules/character.rs
+++ b/src/modules/character.rs
@@ -10,7 +10,7 @@ use clap::ArgMatches;
 /// (green by default)
 /// - If the exit-code was anything else, the arrow will be formatted with
 /// `COLOR_FAILURE` (red by default)
-pub fn segment(args: &ArgMatches) -> Segment {
+pub fn segment(args: &ArgMatches) -> Option<Segment> {
     const PROMPT_CHAR: &str = "➜";
     const COLOR_SUCCESS: Color = Color::Green;
     const COLOR_FAILURE: Color = Color::Red;
@@ -23,7 +23,9 @@ pub fn segment(args: &ArgMatches) -> Segment {
         segment.set_style(COLOR_FAILURE);
     };
 
-    segment.set_value(PROMPT_CHAR).clone()
+    segment.set_value(PROMPT_CHAR).set_prefix(None);
+
+    Some(segment)
 }
 
 #[cfg(test)]
diff --git a/src/modules/directory.rs b/src/modules/directory.rs
index 6ea6f71c3..d62a41f67 100644
--- a/src/modules/directory.rs
+++ b/src/modules/directory.rs
@@ -15,7 +15,7 @@ use std::path::Path;
 ///
 /// **Truncation**
 /// Paths will be limited in length to `3` path components by default.
-pub fn segment(_: &ArgMatches) -> Segment {
+pub fn segment(_: &ArgMatches) -> Option<Segment> {
     const HOME_SYMBOL: &str = "~";
     const DIR_TRUNCATION_LENGTH: usize = 3;
     const SECTION_COLOR: Color = Color::Cyan;
@@ -44,8 +44,9 @@ pub fn segment(_: &ArgMatches) -> Segment {
 
     segment
         .set_value(truncated_dir_string)
-        .set_style(SECTION_COLOR.bold())
-        .clone()
+        .set_style(SECTION_COLOR.bold());
+
+    Some(segment)
 }
 
 /// Get the root directory of a git repo
diff --git a/src/modules/line_break.rs b/src/modules/line_break.rs
index c0c79ea4c..ed8b5023d 100644
--- a/src/modules/line_break.rs
+++ b/src/modules/line_break.rs
@@ -2,10 +2,15 @@ use super::Segment;
 use clap::ArgMatches;
 
 /// Creates a segment for the line break
-pub fn segment(_: &ArgMatches) -> Segment {
+pub fn segment(_: &ArgMatches) -> Option<Segment> {
     const LINE_ENDING: &str = "\n";
 
     let mut segment = Segment::new("line_break");
 
-    segment.set_value(LINE_ENDING).clone()
+    segment
+        .set_value(LINE_ENDING)
+        .set_prefix(None)
+        .set_suffix(None);
+
+    Some(segment)
 }
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index 42620db9f..6fc14d1bf 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -9,7 +9,7 @@ use clap::ArgMatches;
 // pub static current_dir: PathBuf = env::current_dir().expect("Unable to identify current directory");
 // TODO: Currently gets the physical directory. Get the logical directory.
 
-pub fn handle(module: &str, args: &ArgMatches) -> Segment {
+pub fn handle(module: &str, args: &ArgMatches) -> Option<Segment> {
     match module {
         "dir" | "directory" => directory::segment(&args),
         "char" | "character" => character::segment(&args),
diff --git a/src/modules/nodejs.rs b/src/modules/nodejs.rs
index 67032b407..e7a3e736f 100644
--- a/src/modules/nodejs.rs
+++ b/src/modules/nodejs.rs
@@ -11,7 +11,7 @@ use std::process::Command;
 ///     - Current directory contains a `.js` file
 ///     - Current directory contains a `node_modules` directory
 ///     - Current directory contains a `package.json` file
-pub fn segment(_: &ArgMatches) -> Segment {
+pub fn segment(_: &ArgMatches) -> Option<Segment> {
     const NODE_CHAR: &str = "⬢";
     const SECTION_COLOR: Color = Color::Green;
 
@@ -22,7 +22,7 @@ pub fn segment(_: &ArgMatches) -> Segment {
     // Early return if there are no JS project files
     let is_js_project = files.filter_map(Result::ok).any(has_js_files);
     if !is_js_project {
-        return segment;
+        return None;
     }
 
     match Command::new("node").arg("--version").output() {
@@ -31,11 +31,12 @@ pub fn segment(_: &ArgMatches) -> Segment {
             segment.set_value(format!("{} {}", NODE_CHAR, version.trim()))
         }
         Err(_) => {
-            return segment;
+            return None;
         }
     };
 
-    segment.set_style(SECTION_COLOR).clone()
+    segment.set_style(SECTION_COLOR);
+    Some(segment)
 }
 
 fn has_js_files(dir_entry: DirEntry) -> bool {
diff --git a/src/print.rs b/src/print.rs
index 19b13ec4d..cc1719199 100644
--- a/src/print.rs
+++ b/src/print.rs
@@ -18,7 +18,9 @@ pub fn prompt(args: ArgMatches) {
 
     default_prompt
         .iter()
-        .map(|module| modules::handle(module, &args))
-        .map(|segment| segment.output())
+        .map(|module| modules::handle(module, &args)) // Compute segments
+        .flatten() // Remove segments set to `None`
+        .enumerate() // Turn segment into tuple with index
+        .map(|(index, segment)| segment.output(index)) // Generate string outputs
         .for_each(|segment_string| write!(handle, "{}", segment_string).unwrap());
 }
diff --git a/src/segment.rs b/src/segment.rs
index 1e42deb5f..3829c1593 100644
--- a/src/segment.rs
+++ b/src/segment.rs
@@ -56,13 +56,13 @@ impl Segment {
         self
     }
 
-    pub fn set_prefix(&mut self, prefix: Segment) -> &mut Segment {
-        self.prefix = Some(Box::new(prefix));
+    pub fn set_prefix(&mut self, prefix: OptionalSegment) -> &mut Segment {
+        self.prefix = prefix;
         self
     }
 
-    pub fn set_suffix(&mut self, suffix: Segment) -> &mut Segment {
-        self.suffix = Some(Box::new(suffix));
+    pub fn set_suffix(&mut self, suffix: OptionalSegment) -> &mut Segment {
+        self.suffix = suffix;
         self
     }
 
@@ -70,7 +70,7 @@ impl Segment {
     ///
     /// Will recursively also format the prefix and suffix of the segment being
     /// stringified.
-    pub fn output<'a>(&'a self) -> String {
+    pub fn output(&self, index: usize) -> String {
         let Segment {
             name: _name,
             prefix,
@@ -81,14 +81,17 @@ impl Segment {
 
         let mut segment_string = String::new();
 
-        if let Some(prefix) = prefix {
-            segment_string += &prefix.output()
+        // Skip the prefix for the first segment
+        if index != 0 {
+            if let Some(prefix) = prefix {
+                segment_string += &prefix.output(index)
+            }
         }
 
         segment_string += &style.paint(value).to_string();
 
         if let Some(suffix) = suffix {
-            segment_string += &suffix.output();
+            segment_string += &suffix.output(index);
         }
 
         segment_string