From c6ee5c6ac16d360ab1a44d097c91fe9f98f20f85 Mon Sep 17 00:00:00 2001
From: Matan Kushner <hello@matchai.me>
Date: Wed, 1 May 2019 16:34:24 -0400
Subject: [PATCH] Refactor segments into modules (#40)

---
 src/lib.rs                |   1 +
 src/main.rs               |   1 +
 src/module.rs             | 159 +++++++++++++++++++++++++++++++
 src/modules/character.rs  |  25 +++--
 src/modules/directory.rs  |  18 ++--
 src/modules/git_branch.rs |  26 ++---
 src/modules/line_break.rs |  17 ++--
 src/modules/mod.rs        |   4 +-
 src/modules/nodejs.rs     |  18 ++--
 src/modules/package.rs    |  21 +++--
 src/modules/python.rs     |  19 ++--
 src/modules/rust.rs       |  19 ++--
 src/print.rs              |  22 +++--
 src/segment.rs            | 193 ++++++++++++++++++--------------------
 tests/character.rs        |  21 ++---
 tests/common.rs           |  10 +-
 tests/directory.rs        |  91 +++++++++---------
 tests/nodejs.rs           |  36 ++++---
 18 files changed, 423 insertions(+), 278 deletions(-)
 create mode 100644 src/module.rs

diff --git a/src/lib.rs b/src/lib.rs
index 3028646c..078fcd02 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,5 +1,6 @@
 // Lib is present to allow for benchmarking
 pub mod context;
+pub mod module;
 pub mod modules;
 pub mod print;
 pub mod segment;
diff --git a/src/main.rs b/src/main.rs
index 233f7217..d5422a12 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,6 +6,7 @@ extern crate dirs;
 extern crate git2;
 
 mod context;
+mod module;
 mod modules;
 mod print;
 mod segment;
diff --git a/src/module.rs b/src/module.rs
new file mode 100644
index 00000000..903d5cf5
--- /dev/null
+++ b/src/module.rs
@@ -0,0 +1,159 @@
+use crate::segment::Segment;
+use ansi_term::Style;
+use ansi_term::{ANSIString, ANSIStrings};
+use std::fmt;
+use std::string::ToString;
+
+/// A module is a collection of segments showing data for a single integration
+/// (e.g. The git module shows the current git branch and status)
+pub struct Module {
+    /// The module's name, to be used in configuration and logging.
+    name: String,
+
+    /// The styling to be inherited by all segments contained within this module.
+    style: Style,
+
+    /// The prefix used to separate the current module from the previous one.
+    prefix: ModuleAffix,
+
+    /// The collection of segments that compose this module.
+    segments: Vec<Segment>,
+
+    /// The suffix used to separate the current module from the next one.
+    suffix: ModuleAffix,
+}
+
+impl Module {
+    /// Creates a module with no segments.
+    pub fn new(name: &str) -> Module {
+        Module {
+            name: name.to_string(),
+            style: Style::default(),
+            prefix: ModuleAffix::default_prefix(name.to_string()),
+            segments: Vec::new(),
+            suffix: ModuleAffix::default_suffix(name.to_string()),
+        }
+    }
+
+    /// Get a reference to a newly created segment in the module
+    pub fn new_segment<T>(&mut self, name: &str, value: T) -> &mut Segment
+    where
+        T: Into<String>,
+    {
+        let mut segment = Segment::new(name);
+        segment.set_style(self.style);
+        segment.set_value(value.into());
+        self.segments.push(segment);
+
+        self.segments.last_mut().unwrap()
+    }
+
+    /// Get the module's prefix
+    pub fn get_prefix(&mut self) -> &mut ModuleAffix {
+        &mut self.prefix
+    }
+
+    /// Get the module's suffix
+    pub fn get_suffix(&mut self) -> &mut ModuleAffix {
+        &mut self.suffix
+    }
+
+    /// Sets the style of the segment.
+    ///
+    /// Accepts either `Color` or `Style`.
+    pub fn set_style<T>(&mut self, style: T) -> &mut Module
+    where
+        T: Into<Style>,
+    {
+        self.style = style.into();
+        self
+    }
+
+    /// Returns a vector of colored ANSIString elements to be later used with
+    /// `ANSIStrings()` to optimize ANSI codes
+    pub fn ansi_strings(&self) -> Vec<ANSIString> {
+        let mut ansi_strings = self
+            .segments
+            .iter()
+            .map(|s| s.ansi_strings())
+            .flat_map(|s| s.into_iter())
+            .collect::<Vec<ANSIString>>();
+
+        ansi_strings.insert(0, self.prefix.ansi_string());
+        ansi_strings.push(self.suffix.ansi_string());
+
+        ansi_strings
+    }
+
+    pub fn to_string_without_prefix(&self) -> String {
+        ANSIStrings(&self.ansi_strings()[1..]).to_string()
+    }
+}
+
+impl fmt::Display for Module {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let ansi_strings = self.ansi_strings();
+        write!(f, "{}", ANSIStrings(&ansi_strings))
+    }
+}
+
+/// Module affixes are to be used for the prefix or suffix of a module.
+pub struct ModuleAffix {
+    /// The affix's name, to be used in configuration and logging.
+    name: String,
+
+    /// The affix's style.
+    style: Style,
+
+    /// The string value of the affix.
+    value: String,
+}
+
+impl ModuleAffix {
+    pub fn default_prefix(name: String) -> ModuleAffix {
+        ModuleAffix {
+            name: format!("{}_prefix", name),
+            style: Style::default(),
+            value: "via ".to_string(),
+        }
+    }
+
+    pub fn default_suffix(name: String) -> ModuleAffix {
+        ModuleAffix {
+            name: format!("{}_suffix", name),
+            style: Style::default(),
+            value: " ".to_string(),
+        }
+    }
+
+    /// Sets the style of the module.
+    ///
+    /// Accepts either `Color` or `Style`.
+    pub fn set_style<T>(&mut self, style: T) -> &mut ModuleAffix
+    where
+        T: Into<Style>,
+    {
+        self.style = style.into();
+        self
+    }
+
+    /// Sets the value of the module.
+    pub fn set_value<T>(&mut self, value: T) -> &mut ModuleAffix
+    where
+        T: Into<String>,
+    {
+        self.value = value.into();
+        self
+    }
+
+    /// Generates the colored ANSIString output.
+    pub fn ansi_string(&self) -> ANSIString {
+        self.style.paint(&self.value)
+    }
+}
+
+impl fmt::Display for ModuleAffix {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}", self.ansi_string())
+    }
+}
diff --git a/src/modules/character.rs b/src/modules/character.rs
index f624e19f..2a933489 100644
--- a/src/modules/character.rs
+++ b/src/modules/character.rs
@@ -1,8 +1,6 @@
+use super::{Context, Module};
 use ansi_term::Color;
 
-use super::Segment;
-use crate::context::Context;
-
 /// Creates a segment for the prompt character
 ///
 /// The char segment prints an arrow character in a color dependant on the exit-
@@ -11,21 +9,22 @@ use crate::context::Context;
 /// (green by default)
 /// - If the exit-code was anything else, the arrow will be formatted with
 /// `COLOR_FAILURE` (red by default)
-pub fn segment(context: &Context) -> Option<Segment> {
+pub fn segment(context: &Context) -> Option<Module> {
     const PROMPT_CHAR: &str = "➜";
-    const COLOR_SUCCESS: Color = Color::Green;
-    const COLOR_FAILURE: Color = Color::Red;
+    let color_success = Color::Green.bold();
+    let color_failure = Color::Red.bold();
+
+    let mut module = Module::new("char");
+    module.get_prefix().set_value("");
+
+    let symbol = module.new_segment("symbol", PROMPT_CHAR);
 
-    let mut segment = Segment::new("char");
     let arguments = &context.arguments;
-
     if arguments.value_of("status_code").unwrap() == "0" {
-        segment.set_style(COLOR_SUCCESS);
+        symbol.set_style(color_success.bold());
     } else {
-        segment.set_style(COLOR_FAILURE);
+        symbol.set_style(color_failure.bold());
     };
 
-    segment.set_value(PROMPT_CHAR).set_prefix(None);
-
-    Some(segment)
+    Some(module)
 }
diff --git a/src/modules/directory.rs b/src/modules/directory.rs
index d5fbc88c..a85287a6 100644
--- a/src/modules/directory.rs
+++ b/src/modules/directory.rs
@@ -1,8 +1,7 @@
 use ansi_term::Color;
 use std::path::Path;
 
-use super::Segment;
-use crate::context::Context;
+use super::{Context, Module};
 
 /// Creates a segment with the current directory
 ///
@@ -13,12 +12,14 @@ use crate::context::Context;
 ///
 /// **Truncation**
 /// Paths will be limited in length to `3` path components by default.
-pub fn segment(context: &Context) -> Option<Segment> {
+pub fn segment(context: &Context) -> Option<Module> {
     const HOME_SYMBOL: &str = "~";
     const DIR_TRUNCATION_LENGTH: usize = 3;
-    const SEGMENT_COLOR: Color = Color::Cyan;
+    let module_color = Color::Cyan.bold();
+
+    let mut module = Module::new("directory");
+    module.set_style(module_color);
 
-    let mut segment = Segment::new("dir");
     let current_dir = &context.current_dir;
 
     let dir_string;
@@ -37,12 +38,9 @@ pub fn segment(context: &Context) -> Option<Segment> {
 
     // Truncate the dir string to the maximum number of path components
     let truncated_dir_string = truncate(dir_string, DIR_TRUNCATION_LENGTH);
+    module.new_segment("path", truncated_dir_string);
 
-    segment
-        .set_value(truncated_dir_string)
-        .set_style(SEGMENT_COLOR.bold());
-
-    Some(segment)
+    Some(module)
 }
 
 /// Contract the root component of a path
diff --git a/src/modules/git_branch.rs b/src/modules/git_branch.rs
index 586724b2..a89e77f3 100644
--- a/src/modules/git_branch.rs
+++ b/src/modules/git_branch.rs
@@ -1,13 +1,12 @@
 use ansi_term::Color;
 use git2::Repository;
 
-use super::Segment;
-use crate::context::Context;
+use super::{Context, Module};
 
 /// Creates a segment with the Git branch in the current directory
 ///
 /// Will display the branch name if the current directory is a git repo
-pub fn segment(context: &Context) -> Option<Segment> {
+pub fn segment(context: &Context) -> Option<Module> {
     if context.repository.is_none() {
         return None;
     }
@@ -15,22 +14,17 @@ pub fn segment(context: &Context) -> Option<Segment> {
     let repository = context.repository.as_ref().unwrap();
     match get_current_branch(repository) {
         Ok(branch_name) => {
-            const GIT_BRANCH_CHAR: &str = "";
-            const SEGMENT_COLOR: Color = Color::Purple;
+            const GIT_BRANCH_CHAR: &str = " ";
+            let segment_color = Color::Purple.bold();
 
-            // TODO: Make the prefix for the module "in "
-            let mut segment_prefix = Segment::new("git_branch_prefix");
-            segment_prefix
-                .set_value(GIT_BRANCH_CHAR)
-                .set_style(SEGMENT_COLOR.bold());
+            let mut module = Module::new("git_branch");
+            module.set_style(segment_color);
+            module.get_prefix().set_value("in ");
 
-            let mut segment = Segment::new("git_branch");
-            segment
-                .set_prefix(Some(Box::new(segment_prefix)))
-                .set_style(SEGMENT_COLOR.bold())
-                .set_value(branch_name);
+            module.new_segment("branch_char", GIT_BRANCH_CHAR);
+            module.new_segment("branch_name", branch_name);
 
-            Some(segment)
+            Some(module)
         }
         Err(_e) => None,
     }
diff --git a/src/modules/line_break.rs b/src/modules/line_break.rs
index de0e6f9b..223695dc 100644
--- a/src/modules/line_break.rs
+++ b/src/modules/line_break.rs
@@ -1,16 +1,15 @@
-use super::Segment;
-use crate::context::Context;
+use super::{Context, Module};
 
 /// Creates a segment for the line break
-pub fn segment(_context: &Context) -> Option<Segment> {
+pub fn segment(_context: &Context) -> Option<Module> {
     const LINE_ENDING: &str = "\n";
 
-    let mut segment = Segment::new("line_break");
+    let mut module = Module::new("line_break");
 
-    segment
-        .set_value(LINE_ENDING)
-        .set_prefix(None)
-        .set_suffix(None);
+    module.get_prefix().set_value("");
+    module.get_suffix().set_value("");
 
-    Some(segment)
+    module.new_segment("character", LINE_ENDING);
+
+    Some(module)
 }
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index 5c55f6aa..911486c4 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -8,9 +8,9 @@ mod python;
 mod rust;
 
 use crate::context::Context;
-use crate::segment::Segment;
+use crate::module::Module;
 
-pub fn handle(module: &str, context: &Context) -> Option<Segment> {
+pub fn handle(module: &str, context: &Context) -> Option<Module> {
     match module {
         "dir" | "directory" => directory::segment(context),
         "char" | "character" => character::segment(context),
diff --git a/src/modules/nodejs.rs b/src/modules/nodejs.rs
index 818a0cf5..6552c555 100644
--- a/src/modules/nodejs.rs
+++ b/src/modules/nodejs.rs
@@ -2,8 +2,7 @@ use ansi_term::Color;
 use std::path::PathBuf;
 use std::process::Command;
 
-use super::Segment;
-use crate::context::Context;
+use super::{Context, Module};
 
 /// Creates a segment with the current Node.js version
 ///
@@ -11,7 +10,7 @@ use crate::context::Context;
 ///     - Current directory contains a `.js` file
 ///     - Current directory contains a `package.json` file
 ///     - Current directory contains a `node_modules` directory
-pub fn segment(context: &Context) -> Option<Segment> {
+pub fn segment(context: &Context) -> Option<Module> {
     let is_js_project = context.dir_files.iter().any(has_js_files);
     if !is_js_project {
         return None;
@@ -19,16 +18,17 @@ pub fn segment(context: &Context) -> Option<Segment> {
 
     match get_node_version() {
         Some(node_version) => {
-            const NODE_CHAR: &str = "⬢";
-            const SEGMENT_COLOR: Color = Color::Green;
+            const NODE_CHAR: &str = "⬢ ";
+            let module_color = Color::Green.bold();
 
-            let mut segment = Segment::new("node");
-            segment.set_style(SEGMENT_COLOR.bold());
+            let mut module = Module::new("node");
+            module.set_style(module_color);
 
             let formatted_version = node_version.trim();
-            segment.set_value(format!("{} {}", NODE_CHAR, formatted_version));
+            module.new_segment("symbol", NODE_CHAR);
+            module.new_segment("version", formatted_version);
 
-            Some(segment)
+            Some(module)
         }
         None => None,
     }
diff --git a/src/modules/package.rs b/src/modules/package.rs
index ab201e4b..8e7d4d3a 100644
--- a/src/modules/package.rs
+++ b/src/modules/package.rs
@@ -1,5 +1,5 @@
-use super::Segment;
-use crate::context::Context;
+use super::{Context, Module};
+
 use ansi_term::Color;
 use serde_json;
 use std::fs::File;
@@ -11,19 +11,20 @@ use toml;
 /// Creates a segment with the current package version
 ///
 /// Will display if a version is defined for your Node.js or Rust project (if one exists)
-pub fn segment(context: &Context) -> Option<Segment> {
+pub fn segment(context: &Context) -> Option<Module> {
     match get_package_version(context) {
         Some(package_version) => {
-            const PACKAGE_CHAR: &str = "📦";
-            const SEGMENT_COLOR: Color = Color::Red;
+            const PACKAGE_CHAR: &str = "📦 ";
+            let module_color = Color::Red.bold();
 
-            // TODO: Make the prefix for the module "is "
-            let mut segment = Segment::new("package");
-            segment.set_style(SEGMENT_COLOR.bold());
+            let mut module = Module::new("package");
+            module.set_style(module_color);
+            module.get_prefix().set_value("is ");
 
-            segment.set_value(format!("{} {}", PACKAGE_CHAR, package_version));
+            module.new_segment("symbol", PACKAGE_CHAR);
+            module.new_segment("version", package_version);
 
-            Some(segment)
+            Some(module)
         }
         None => None,
     }
diff --git a/src/modules/python.rs b/src/modules/python.rs
index afd62690..70625a47 100644
--- a/src/modules/python.rs
+++ b/src/modules/python.rs
@@ -1,9 +1,9 @@
-use super::Segment;
-use crate::context::Context;
 use ansi_term::Color;
 use std::path::PathBuf;
 use std::process::Command;
 
+use super::{Context, Module};
+
 /// Creates a segment with the current Python version
 ///
 /// Will display the Python version if any of the following criteria are met:
@@ -11,7 +11,7 @@ use std::process::Command;
 ///     - Current directory contains a `.python-version` file
 ///     - Current directory contains a `requirements.txt` file
 ///     - Current directory contains a `pyproject.toml` file
-pub fn segment(context: &Context) -> Option<Segment> {
+pub fn segment(context: &Context) -> Option<Module> {
     let is_py_project = context.dir_files.iter().any(has_py_files);
     if !is_py_project {
         return None;
@@ -19,16 +19,17 @@ pub fn segment(context: &Context) -> Option<Segment> {
 
     match get_python_version() {
         Some(python_version) => {
-            const PYTHON_CHAR: &str = "🐍";
-            const SEGMENT_COLOR: Color = Color::Yellow;
+            const PYTHON_CHAR: &str = "🐍 ";
+            let module_color = Color::Yellow.bold();
 
-            let mut segment = Segment::new("python");
-            segment.set_style(SEGMENT_COLOR.bold());
+            let mut module = Module::new("python");
+            module.set_style(module_color);
 
             let formatted_version = format_python_version(python_version);
-            segment.set_value(format!("{} {}", PYTHON_CHAR, formatted_version));
+            module.new_segment("symbol", PYTHON_CHAR);
+            module.new_segment("version", formatted_version);
 
-            Some(segment)
+            Some(module)
         }
         None => None,
     }
diff --git a/src/modules/rust.rs b/src/modules/rust.rs
index 14af3b85..1f0ae603 100644
--- a/src/modules/rust.rs
+++ b/src/modules/rust.rs
@@ -1,15 +1,15 @@
-use super::Segment;
-use crate::context::Context;
 use ansi_term::Color;
 use std::path::PathBuf;
 use std::process::Command;
 
+use super::{Context, Module};
+
 /// Creates a segment with the current Rust version
 ///
 /// Will display the Rust version if any of the following criteria are met:
 ///     - Current directory contains a `.rs` file
 ///     - Current directory contains a `Cargo.toml` file
-pub fn segment(context: &Context) -> Option<Segment> {
+pub fn segment(context: &Context) -> Option<Module> {
     let is_rs_project = context.dir_files.iter().any(has_rs_files);
     if !is_rs_project {
         return None;
@@ -17,16 +17,17 @@ pub fn segment(context: &Context) -> Option<Segment> {
 
     match get_rust_version() {
         Some(rust_version) => {
-            const RUST_CHAR: &str = "🦀";
-            const SEGMENT_COLOR: Color = Color::Red;
+            const RUST_CHAR: &str = "🦀 ";
+            let module_color = Color::Red.bold();
 
-            let mut segment = Segment::new("rust");
-            segment.set_style(SEGMENT_COLOR.bold());
+            let mut module = Module::new("rust");
+            module.set_style(module_color);
 
             let formatted_version = format_rustc_version(rust_version);
-            segment.set_value(format!("{} {}", RUST_CHAR, formatted_version));
+            module.new_segment("symbol", RUST_CHAR);
+            module.new_segment("version", formatted_version);
 
-            Some(segment)
+            Some(module)
         }
         None => None,
     }
diff --git a/src/print.rs b/src/print.rs
index e2f0a066..780d9099 100644
--- a/src/print.rs
+++ b/src/print.rs
@@ -2,6 +2,7 @@ use clap::ArgMatches;
 use std::io::{self, Write};
 
 use crate::context::Context;
+use crate::module::Module;
 use crate::modules;
 
 pub fn prompt(args: ArgMatches) {
@@ -27,11 +28,20 @@ pub fn prompt(args: ArgMatches) {
     // Write a new line before the prompt
     writeln!(handle).unwrap();
 
-    prompt_order
+    let modules = prompt_order
         .iter()
-        .map(|module| modules::handle(module, &context)) // Compute segments
-        .flatten() // Remove segments set to `None`
-        .enumerate() // Turn segment into tuple with index
-        .map(|(index, segment)| segment.output_index(index)) // Generate string outputs
-        .for_each(|segment_string| write!(handle, "{}", segment_string).unwrap());
+        .map(|module| modules::handle(module, &context)) // Compute modules
+        .flatten()
+        .collect::<Vec<Module>>(); // Remove segments set to `None`
+
+    let mut printable = modules.iter();
+
+    // Print the first module without its prefix
+    if let Some(first_module) = printable.next() {
+        let module_without_prefix = first_module.to_string_without_prefix();
+        write!(handle, "{}", module_without_prefix).unwrap()
+    }
+
+    // Print all remaining modules
+    printable.for_each(|module| write!(handle, "{}", module).unwrap());
 }
diff --git a/src/segment.rs b/src/segment.rs
index a56e1bf8..0d71af1a 100644
--- a/src/segment.rs
+++ b/src/segment.rs
@@ -1,58 +1,50 @@
-use ansi_term::Style;
+use ansi_term::{ANSIString, ANSIStrings, Style};
+use std::fmt;
 
-#[derive(Clone)]
+/// 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).
 pub struct Segment {
-    name: Option<String>,
-    style: Style,
+    /// The segment's name, to be used in configuration and logging.
+    name: String,
+
+    /// The segment's style. If None, will inherit the style of the module containing it.
+    style: Option<Style>,
+
+    /// The prefix used to preceed the contents of a segment.
+    prefix: Option<SegmentAffix>,
+
+    /// The string value of the current segment.
     value: String,
-    prefix: BoxedSegment,
-    suffix: BoxedSegment,
+
+    /// The suffix used following the contents of a segment.
+    suffix: Option<SegmentAffix>,
 }
 
 impl Segment {
-    /// Creates a new segment with default fields
-    pub fn new<T>(name: T) -> Segment
-    where
-        T: Into<String>,
-        T: Copy,
-    {
-        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,
-        }));
-
+    /// Creates a new segment with default fields.
+    pub fn new(name: &str) -> Segment {
         Segment {
-            name: Some(name.into()),
-            style: Style::default(),
-            value: String::from(""),
-            prefix: default_prefix,
-            suffix: default_suffix,
+            name: name.to_string(),
+            style: None,
+            prefix: None,
+            value: "".to_string(),
+            suffix: None,
         }
     }
 
-    /// Sets the style of the segment
+    /// Sets the style of the segment.
     ///
     /// Accepts either `Color` or `Style`.
     pub fn set_style<T>(&mut self, style: T) -> &mut Segment
     where
         T: Into<Style>,
     {
-        self.style = style.into();
+        self.style = Some(style.into());
         self
     }
 
-    /// Sets the value of the segment
+    /// Sets the value of the segment.
     pub fn set_value<T>(&mut self, value: T) -> &mut Segment
     where
         T: Into<String>,
@@ -61,77 +53,70 @@ impl Segment {
         self
     }
 
-    /// Sets the prefix of the segment
-    pub fn set_prefix(&mut self, prefix: BoxedSegment) -> &mut Segment {
-        self.prefix = prefix;
-        self
+    // Returns the ANSIString of the segment value, not including its prefix and suffix
+    fn value_ansi_string(&self) -> ANSIString {
+        match self.style {
+            Some(style) => style.paint(&self.value),
+            None => ANSIString::from(&self.value),
+        }
     }
 
-    /// Sets the suffix of the segment
-    pub fn set_suffix(&mut self, suffix: BoxedSegment) -> &mut Segment {
-        self.suffix = suffix;
-        self
-    }
+    /// Returns a vector of colored ANSIString elements to be later used with
+    /// `ANSIStrings()` to optimize ANSI codes
+    pub fn ansi_strings(&self) -> Vec<ANSIString> {
+        let prefix = self.prefix.as_ref().and_then(|p| Some(p.ansi_string()));
+        let suffix = self.suffix.as_ref().and_then(|s| Some(s.ansi_string()));
+        let value = Some(self.value_ansi_string());
 
-    /// Create a string with the formatted contents of a segment
-    ///
-    /// Will recursively also format the prefix and suffix of the segment being
-    /// stringified.
-    pub fn output(&self) -> String {
-        let Segment {
-            name: _name,
-            prefix,
-            value,
-            style,
-            suffix,
-        } = self;
-
-        let mut segment_string = String::new();
-
-        // Skip the prefix for the first segment
-        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
-    }
-
-    /// Create a string with the formatted contents of a segment while skipping the first segment.
-    ///
-    /// Will recursively also format the prefix and suffix of the segment being
-    /// stringified.
-    pub fn output_index(&self, index: usize) -> String {
-        let Segment {
-            name: _name,
-            prefix,
-            value,
-            style,
-            suffix,
-        } = self;
-
-        let mut segment_string = String::new();
-
-        // Skip the prefix for the first segment
-        if index != 0 {
-            if let Some(prefix) = prefix {
-                segment_string += &prefix.output_index(index)
-            }
-        }
-
-        segment_string += &style.paint(value).to_string();
-
-        if let Some(suffix) = suffix {
-            segment_string += &suffix.output();
-        }
-
-        segment_string
+        // Remove `None` values from the vector
+        vec![prefix, value, suffix]
+            .into_iter()
+            .filter_map(|e| e)
+            .collect::<Vec<ANSIString>>()
     }
 }
 
-type BoxedSegment = Option<Box<Segment>>;
+impl fmt::Display for Segment {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let ansi_strings = self.ansi_strings();
+        write!(f, "{}", ANSIStrings(&ansi_strings))
+    }
+}
+
+/// Segment affixes are to be used for the prefix or suffix of a segment.
+/// By default they will inherit the styling of its segment, unless otherwise specified.
+pub struct SegmentAffix {
+    /// The affix's name, to be used in configuration and logging.
+    name: String,
+
+    /// The affix's style. If None, will inherit the style of the segment containing it.
+    style: Option<Style>,
+
+    /// The string value of the affix.
+    value: String,
+}
+
+impl SegmentAffix {
+    /// Creates a segment affix with no contents.
+    pub fn new() -> SegmentAffix {
+        SegmentAffix {
+            name: String::new(),
+            style: None,
+            value: String::new(),
+        }
+    }
+
+    /// Generates the colored ANSIString output.
+    pub fn ansi_string(&self) -> ANSIString {
+        match self.style {
+            Some(style) => style.paint(&self.value),
+            None => ANSIString::from(&self.value),
+        }
+    }
+}
+
+impl fmt::Display for SegmentAffix {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}", self.ansi_string())
+    }
+}
diff --git a/tests/character.rs b/tests/character.rs
index 511423da..be0b1a8e 100644
--- a/tests/character.rs
+++ b/tests/character.rs
@@ -1,29 +1,20 @@
 use ansi_term::Color;
-use starship::segment::Segment;
 use std::path::Path;
 
 mod common;
 
 #[test]
-fn char_segment_success_status() {
+fn char_module_success_status() {
     let dir = Path::new("~");
-    let expected = Segment::new("char")
-        .set_value("➜")
-        .set_style(Color::Green)
-        .set_prefix(None)
-        .output();
-    let actual = common::render_segment_with_status("char", &dir, "0");
+    let expected = format!("{} ", Color::Green.bold().paint("➜"));
+    let actual = common::render_module_with_status("char", &dir, "0");
     assert_eq!(expected, actual);
 }
 
 #[test]
-fn char_segment_failure_status() {
+fn char_module_failure_status() {
     let dir = Path::new("~");
-    let expected = Segment::new("char")
-        .set_value("➜")
-        .set_style(Color::Red)
-        .set_prefix(None)
-        .output();
-    let actual = common::render_segment_with_status("char", &dir, "1");
+    let expected = format!("{} ", Color::Red.bold().paint("➜"));
+    let actual = common::render_module_with_status("char", &dir, "1");
     assert_eq!(expected, actual);
 }
diff --git a/tests/common.rs b/tests/common.rs
index 040732ff..2d2cd2a0 100644
--- a/tests/common.rs
+++ b/tests/common.rs
@@ -4,14 +4,14 @@ use starship::modules;
 use std::path::PathBuf;
 
 #[allow(dead_code)]
-pub fn render_segment<T>(module: &str, path: T) -> String
+pub fn render_module<T>(module: &str, path: T) -> String
 where
     T: Into<PathBuf>,
 {
-    render_segment_with_status(module, path.into(), "0")
+    render_module_with_status(module, path.into(), "0")
 }
 
-pub fn render_segment_with_status<T>(module: &str, path: T, status: &str) -> String
+pub fn render_module_with_status<T>(module: &str, path: T, status: &str) -> String
 where
     T: Into<PathBuf>,
 {
@@ -21,7 +21,7 @@ where
         .get_matches_from(vec!["starship", status]);
     let context = Context::new_with_dir(args, path.into());
 
-    let segment = modules::handle(module, &context);
+    let module = modules::handle(module, &context);
 
-    segment.unwrap().output()
+    module.unwrap().to_string()
 }
diff --git a/tests/directory.rs b/tests/directory.rs
index 2eb20f48..38efc298 100644
--- a/tests/directory.rs
+++ b/tests/directory.rs
@@ -1,7 +1,6 @@
 use ansi_term::Color;
 use dirs::home_dir;
 use git2::Repository;
-use starship::segment::Segment;
 use std::fs;
 use std::io;
 use std::path::Path;
@@ -13,11 +12,8 @@ mod common;
 fn home_directory() -> io::Result<()> {
     let dir = Path::new("~");
 
-    let expected = Segment::new("dir")
-        .set_value("~")
-        .set_style(Color::Cyan.bold())
-        .output();
-    let actual = common::render_segment("dir", &dir);
+    let expected = format!("via {} ", Color::Cyan.bold().paint("~").to_string());
+    let actual = common::render_module("dir", &dir);
     assert_eq!(expected, actual);
 
     Ok(())
@@ -29,11 +25,11 @@ fn directory_in_home() -> io::Result<()> {
     let dir = home_dir().unwrap().join("starship/engine");
     fs::create_dir_all(&dir)?;
 
-    let expected = Segment::new("dir")
-        .set_value("~/starship/engine")
-        .set_style(Color::Cyan.bold())
-        .output();
-    let actual = common::render_segment("dir", &dir);
+    let expected = format!(
+        "via {} ",
+        Color::Cyan.bold().paint("~/starship/engine").to_string()
+    );
+    let actual = common::render_module("dir", &dir);
     assert_eq!(expected, actual);
 
     Ok(())
@@ -45,11 +41,14 @@ fn truncated_directory_in_home() -> io::Result<()> {
     let dir = home_dir().unwrap().join("starship/engine/schematics");
     fs::create_dir_all(&dir)?;
 
-    let expected = Segment::new("dir")
-        .set_value("starship/engine/schematics")
-        .set_style(Color::Cyan.bold())
-        .output();
-    let actual = common::render_segment("dir", &dir);
+    let expected = format!(
+        "via {} ",
+        Color::Cyan
+            .bold()
+            .paint("starship/engine/schematics")
+            .to_string()
+    );
+    let actual = common::render_module("dir", &dir);
     assert_eq!(expected, actual);
 
     Ok(())
@@ -59,11 +58,8 @@ fn truncated_directory_in_home() -> io::Result<()> {
 fn root_directory() -> io::Result<()> {
     let dir = Path::new("/");
 
-    let expected = Segment::new("dir")
-        .set_value("/")
-        .set_style(Color::Cyan.bold())
-        .output();
-    let actual = common::render_segment("dir", &dir);
+    let expected = format!("via {} ", Color::Cyan.bold().paint("/").to_string());
+    let actual = common::render_module("dir", &dir);
     assert_eq!(expected, actual);
 
     Ok(())
@@ -73,11 +69,8 @@ fn root_directory() -> io::Result<()> {
 fn directory_in_root() -> io::Result<()> {
     let dir = Path::new("/opt");
 
-    let expected = Segment::new("dir")
-        .set_value("/opt")
-        .set_style(Color::Cyan.bold())
-        .output();
-    let actual = common::render_segment("dir", &dir);
+    let expected = format!("via {} ", Color::Cyan.bold().paint("/opt").to_string());
+    let actual = common::render_module("dir", &dir);
     assert_eq!(expected, actual);
 
     Ok(())
@@ -89,11 +82,14 @@ fn truncated_directory_in_root() -> io::Result<()> {
     let dir = Path::new("/opt/starship/thrusters/rocket");
     fs::create_dir_all(&dir)?;
 
-    let expected = Segment::new("dir")
-        .set_value("starship/thrusters/rocket")
-        .set_style(Color::Cyan.bold())
-        .output();
-    let actual = common::render_segment("dir", &dir);
+    let expected = format!(
+        "via {} ",
+        Color::Cyan
+            .bold()
+            .paint("starship/thrusters/rocket")
+            .to_string()
+    );
+    let actual = common::render_module("dir", &dir);
     assert_eq!(expected, actual);
 
     Ok(())
@@ -108,11 +104,11 @@ fn git_repo_root() -> io::Result<()> {
 
     Repository::init(&repo_dir).unwrap();
 
-    let expected = Segment::new("dir")
-        .set_value("rocket-controls")
-        .set_style(Color::Cyan.bold())
-        .output();
-    let actual = common::render_segment("dir", &repo_dir);
+    let expected = format!(
+        "via {} ",
+        Color::Cyan.bold().paint("rocket-controls").to_string()
+    );
+    let actual = common::render_module("dir", &repo_dir);
     assert_eq!(expected, actual);
 
     Ok(())
@@ -128,11 +124,11 @@ fn directory_in_git_repo() -> io::Result<()> {
 
     Repository::init(&repo_dir).unwrap();
 
-    let expected = Segment::new("dir")
-        .set_value("rocket-controls/src")
-        .set_style(Color::Cyan.bold())
-        .output();
-    let actual = common::render_segment("dir", &dir);
+    let expected = format!(
+        "via {} ",
+        Color::Cyan.bold().paint("rocket-controls/src").to_string()
+    );
+    let actual = common::render_module("dir", &dir);
     assert_eq!(expected, actual);
 
     Ok(())
@@ -148,11 +144,14 @@ fn truncated_directory_in_git_repo() -> io::Result<()> {
 
     Repository::init(&repo_dir).unwrap();
 
-    let expected = Segment::new("dir")
-        .set_value("src/meters/fuel-gauge")
-        .set_style(Color::Cyan.bold())
-        .output();
-    let actual = common::render_segment("dir", &dir);
+    let expected = format!(
+        "via {} ",
+        Color::Cyan
+            .bold()
+            .paint("src/meters/fuel-gauge")
+            .to_string()
+    );
+    let actual = common::render_module("dir", &dir);
     assert_eq!(expected, actual);
 
     Ok(())
diff --git a/tests/nodejs.rs b/tests/nodejs.rs
index fe07a0cd..38bca73b 100644
--- a/tests/nodejs.rs
+++ b/tests/nodejs.rs
@@ -12,11 +12,13 @@ fn folder_with_package_json() -> io::Result<()> {
     let dir = TempDir::new()?;
     File::create(dir.path().join("package.json"))?;
 
-    let expected = Segment::new("node")
-        .set_value("⬢ v12.0.0")
-        .set_style(Color::Green.bold())
-        .output();
-    let actual = common::render_segment("nodejs", &dir.path());
+    let expected = format!(
+        "via {} ",
+        Segment::new("node")
+            .set_value("⬢ v12.0.0")
+            .set_style(Color::Green.bold())
+    );
+    let actual = common::render_module("nodejs", &dir.path());
     assert_eq!(expected, actual);
 
     Ok(())
@@ -28,11 +30,13 @@ fn folder_with_js_file() -> io::Result<()> {
     let dir = TempDir::new()?;
     File::create(dir.path().join("index.js"))?;
 
-    let expected = Segment::new("node")
-        .set_value("⬢ v12.0.0")
-        .set_style(Color::Green.bold())
-        .output();
-    let actual = common::render_segment("nodejs", &dir.path());
+    let expected = format!(
+        "via {} ",
+        Segment::new("node")
+            .set_value("⬢ v12.0.0")
+            .set_style(Color::Green.bold())
+    );
+    let actual = common::render_module("nodejs", &dir.path());
     assert_eq!(expected, actual);
 
     Ok(())
@@ -45,11 +49,13 @@ fn folder_with_node_modules() -> io::Result<()> {
     let node_modules = dir.path().join("node_modules");
     fs::create_dir_all(&node_modules)?;
 
-    let expected = Segment::new("node")
-        .set_value("⬢ v12.0.0")
-        .set_style(Color::Green.bold())
-        .output();
-    let actual = common::render_segment("nodejs", &dir.path());
+    let expected = format!(
+        "via {} ",
+        Segment::new("node")
+            .set_value("⬢ v12.0.0")
+            .set_style(Color::Green.bold())
+    );
+    let actual = common::render_module("nodejs", &dir.path());
     assert_eq!(expected, actual);
 
     Ok(())