Add dir section with home truncation

4 files changed, 62 insertions, 3 deletions

diff --git a/src/modules/char.rs b/src/modules/char.rs
index 93d8df692..af24ef082 100644
--- a/src/modules/char.rs
+++ b/src/modules/char.rs
@@ -2,7 +2,7 @@ use super::Segment;
 use ansi_term::{Color, Style};
 use clap::ArgMatches;
 
-/// Prints the prompt character
+/// Creates a segment for the prompt character
 /// 
 /// The char segment prints an arrow character in a color dependant on the exit-
 /// code of the last executed command:
diff --git a/src/modules/dir.rs b/src/modules/dir.rs
new file mode 100644
index 000000000..5f17d709b
--- /dev/null
+++ b/src/modules/dir.rs
@@ -0,0 +1,50 @@
+use super::Segment;
+use dirs;
+use std::env;
+use ansi_term::{Color, Style};
+use clap::ArgMatches;
+
+/// Creates a segment with the current directory
+pub fn segment(args: &ArgMatches) -> Segment {
+    const COLOR_DIR: Color = Color::Cyan;
+    const HOME_SYMBOL: char = '~';
+
+    let current_dir = env::current_dir().unwrap();
+    let mut dir_string = String::from(current_dir.to_str().unwrap());
+
+    if let Some(home_dir) = dirs::home_dir() {
+        if current_dir.starts_with(&home_dir) {
+            let current_dir = current_dir.to_str().unwrap();
+            let home_dir = home_dir.to_str().unwrap();
+
+            dir_string = current_dir.replace(home_dir, &HOME_SYMBOL.to_string());
+        }
+    }
+
+    Segment {
+        value: dir_string,
+        style: Style::from(COLOR_DIR).bold(),
+        ..Default::default()
+    }
+}
+
+// fn truncate_dir(directory: PathBuf, truncate_to: u8) {
+
+// }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use clap::{App, Arg};
+
+    #[test]
+    fn truncate_home_dir() {
+        let args = App::new("starship")
+            .arg(Arg::with_name("status_code"))
+            .get_matches_from(vec!["starship", "0"]);
+
+        env::set_current_dir("~/dev/");
+        let segment = segment(&args);
+        assert_eq!(segment.style, Style::from(Color::Green));
+    }
+}
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index 869a82c2c..13b18fa13 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -1,4 +1,5 @@
 mod char;
+mod dir;
 
 use clap::ArgMatches;
 use ansi_term::Style;
@@ -12,11 +13,18 @@ pub struct 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: None
+            suffix: default_suffix
         }
     }
 }
@@ -24,6 +32,7 @@ impl Default for Segment {
 pub fn handle(module: &str, args: &ArgMatches) -> Segment {
     match module {
         "char" => char::segment(&args),
+        "dir" => dir::segment(&args),
 
         _ => panic!("Unknown module: {}", module),
     }
diff --git a/src/print.rs b/src/print.rs
index a74655aa4..8eb526698 100644
--- a/src/print.rs
+++ b/src/print.rs
@@ -3,7 +3,7 @@ use crate::modules::Segment;
 use clap::ArgMatches;
 
 pub fn prompt(args: ArgMatches) {
-    let default_prompt = vec!["char"];
+    let default_prompt = vec!["dir", "char"];
 
     for module in default_prompt {
         let segment = modules::handle(module, &args);