Use git project root for truncation

3 files changed, 61 insertions, 12 deletions

diff --git a/src/lib.rs b/src/lib.rs
index ccfc24c86..50c2327d2 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,2 +1,3 @@
+// Lib is present to allow for benchmarking
 pub mod modules;
 pub mod print;
diff --git a/src/main.rs b/src/main.rs
index 2f41fb5dd..9eb013935 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,6 +2,7 @@
 extern crate clap;
 extern crate ansi_term;
 extern crate dirs;
+extern crate git2;
 
 mod modules;
 mod print;
diff --git a/src/modules/directory.rs b/src/modules/directory.rs
index 3e4546e0a..5f4bb5f53 100644
--- a/src/modules/directory.rs
+++ b/src/modules/directory.rs
@@ -1,33 +1,80 @@
+use std::env;
+use std::path::{Path, PathBuf};
 use super::Segment;
+use git2::{Repository};
 use ansi_term::{Color, Style};
 use clap::ArgMatches;
 use dirs;
-use std::env;
 
 /// Creates a segment with the current directory
 pub fn segment(_: &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());
+    let current_dir = env::current_dir().expect("Unable to identify current directory");
 
-    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());
-        }
+    let dirname;
+    if let Ok(repo) = git2::Repository::discover(&current_dir) {
+        let repo_root = get_repo_root(repo);
+        let repo_root_basename = repo_root.components().last().unwrap();
+        let basename_str_slice = repo_root_basename.as_os_str();
+        dirname = basename_str_slice.to_str().unwrap().to_string();
+    } else {
+        dirname = String::from("test");
     }
 
+    // let dir_string;
+    // if let Ok(repo) = git2::Repository::discover(&current_dir) {
+    //     let repo_root = get_repo_root(repo);
+    //     let repo_root_basename = repo_root.components().last().unwrap();
+    //     dir_string = *repo_root_basename.as_os_str().to_str().unwrap();
+    // } else {
+    //     dir_string = match truncate_home(current_dir) {
+    //         Some(dir) => &dir.to_string(),
+    //         None => &current_dir.to_str().unwrap()
+    //     }
+    // }
+
+    // if let Love(tiff) = matan::Love(tiff) {
+    //     log tiff + matan + kimu + nimu + puku + owl fren + roomba fren + cactus fren + rumple
+    // }
+
+    // let mut dir_string = String::from(current_dir.to_str().unwrap());
+
     Segment {
-        value: dir_string,
+        value: String::from(dirname),
         style: Style::from(COLOR_DIR).bold(),
         ..Default::default()
     }
 }
 
+fn get_repo_root(repo: Repository) -> PathBuf {
+    match repo.is_bare() {
+        // A bare repo will return its root path
+        true => repo.path().to_path_buf(),
+        // Non-bare repos will return the path of `.git`
+        false => repo.path().parent().unwrap().to_path_buf()
+    }
+}
+
+// fn truncate_home(path: PathBuf) -> Option<String> {
+//     const HOME_SYMBOL: &str = "~";
+
+//     if dirs::home_dir() == None {
+//         return None;
+//     }
+
+//     if let Some(home_dir) = dirs::home_dir() {
+//         if path.strip_prefix(home_dir).is_ok() {
+//             let path_str = path.to_str().unwrap();
+//             let home_dir = home_dir.to_str().unwrap();
+            
+//             return Some(path_str.replace(home_dir, HOME_SYMBOL));
+//         }
+//     }
+
+//     None
+// }
+
 #[cfg(test)]
 mod tests {
     // TODO: Look into stubbing `env` so that tests can be run in parallel