7 files changed, 91 insertions, 29 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 50c2327d2..4d3665457 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,3 +1,4 @@
 // Lib is present to allow for benchmarking
 pub mod modules;
 pub mod print;
+pub mod segment;
diff --git a/src/main.rs b/src/main.rs
index b8f11cfc9..f8bc4bcd4 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -7,6 +7,7 @@ extern crate git2;
 
 mod modules;
 mod print;
+mod segment;
 
 use clap::{App, Arg};
 
diff --git a/src/modules/character.rs b/src/modules/character.rs
index 4c83c5096..98b9b254e 100644
--- a/src/modules/character.rs
+++ b/src/modules/character.rs
@@ -15,8 +15,10 @@ pub fn segment(args: &ArgMatches) -> Segment {
     const COLOR_SUCCESS: Color = Color::Green;
     const COLOR_FAILURE: Color = Color::Red;
 
+    let segment = Segment::new("char");
+
     let color = if args.value_of("status_code").unwrap() == "0" {
-        COLOR_SUCCESS
+        segment.set_style(COLOR_SUCCESS);
     } else {
         COLOR_FAILURE
     };
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index 5fa37593e..42620db9f 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -3,37 +3,12 @@ mod directory;
 mod line_break;
 mod nodejs;
 
-use ansi_term::Style;
+use crate::segment::Segment;
 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 struct Segment {
-    pub style: Style,
-    pub value: String,
-    pub prefix: Option<Box<Segment>>,
-    pub suffix: Option<Box<Segment>>,
-}
-
-impl Default for Segment {
-    fn default() -> Segment {
-        let default_suffix = Some(Box::new(Segment {
-            style: Style::default(),
-            value: String::from(" "),
-            prefix: None,
-            suffix: None,
-        }));
-
-        Segment {
-            style: Style::default(),
-            value: String::from(""),
-            prefix: None,
-            suffix: default_suffix,
-        }
-    }
-}
-
 pub fn handle(module: &str, args: &ArgMatches) -> Segment {
     match module {
         "dir" | "directory" => directory::segment(&args),
diff --git a/src/modules/nodejs.rs b/src/modules/nodejs.rs
index 87a1e4a25..a0b4162c0 100644
--- a/src/modules/nodejs.rs
+++ b/src/modules/nodejs.rs
@@ -6,7 +6,7 @@ use std::fs::{self, DirEntry};
 use std::process::Command;
 
 /// Creates a segment with the current Node.js version
-/// 
+///
 /// Will display the Node.js version if any of the following criteria are met:
 ///     - Current directory contains a `.js` file
 ///     - Current directory contains a `node_modules` directory
diff --git a/src/print.rs b/src/print.rs
index 13c367997..03e22374e 100644
--- a/src/print.rs
+++ b/src/print.rs
@@ -5,7 +5,7 @@ use crate::modules;
 use crate::modules::Segment;
 
 pub fn prompt(args: ArgMatches) {
-    let default_prompt = vec!["directory", "node", "line_break", "character"];
+    let default_prompt = vec!["directory", "nodejs", "line_break", "character"];
 
     // TODO:
     // - List files in directory
diff --git a/src/segment.rs b/src/segment.rs
new file mode 100644
index 000000000..df3017628
--- /dev/null
+++ b/src/segment.rs
@@ -0,0 +1,83 @@
+use ansi_term::Style;
+
+pub struct Segment {
+    name: Option<String>,
+    style: Style,
+    value: String,
+    prefix: OptionalSegment,
+    suffix: OptionalSegment,
+}
+
+impl Segment {
+    pub fn new<S>(name: S) -> Segment where S: Into<String> {
+        let default_prefix = Some(Box::new(Segment {
+            name: Some(format!("{} {}", name.into(), "prefix")),
+            style: Style::default(),
+            value: String::from("via "),
+            prefix: None,
+            suffix: None,
+        }));
+
+        let default_suffix = Some(Box::new(Segment {
+            name: Some(format!("{} {}", name.into(), "suffix")),
+            style: Style::default(),
+            value: String::from(" "),
+            prefix: None,
+            suffix: None,
+        }));
+
+        Segment {
+            name: Some(name.into()),
+            style: Style::default(),
+            value: String::from(""),
+            prefix: default_prefix,
+            suffix: default_suffix,
+        }
+    }
+
+    pub fn set_style<'a>(&'a mut self, style: Style) -> &'a mut Segment {
+        self.style = style;
+        self
+    }
+
+    pub fn set_value<'a>(&'a mut self, value: String) -> &'a mut Segment {
+        self.value = value;
+        self
+    }
+
+    pub fn set_prefix<'a>(&'a mut self, prefix: Segment) -> &'a mut Segment {
+        self.prefix = Some(Box::new(prefix));
+        self
+    }
+
+    pub fn set_suffix<'a>(&'a mut self, suffix: Segment) -> &'a mut Segment {
+        self.suffix = Some(Box::new(suffix));
+        self
+    }
+
+    pub fn output<'a>(&'a self) -> String {
+        let Segment {
+            name: _name,
+            prefix,
+            value,
+            style,
+            suffix,
+        } = self;
+
+        let mut segment_string = String::new();
+
+        if let Some(prefix) = prefix {
+            segment_string += &prefix.output()
+        }
+
+        segment_string += &style.paint(value).to_string();
+
+        if let Some(suffix) = suffix {
+            segment_string += &suffix.output();
+        }
+
+        segment_string
+    }
+}
+
+type OptionalSegment = Option<Box<Segment>>;