Add the Git branch segment (#32)

Added
- Repository to Context for reuse in directory and git_branch
- git_branch to prompt

Changed
- Made segments bold to match spaceship

8 files changed, 62 insertions, 20 deletions

diff --git a/src/context.rs b/src/context.rs
index 71662fe06..6baa19536 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -1,4 +1,5 @@
 use clap::ArgMatches;
+use git2::Repository;
 use std::env;
 use std::fs;
 use std::path::PathBuf;
@@ -7,6 +8,7 @@ pub struct Context<'a> {
     pub current_dir: PathBuf,
     pub dir_files: Vec<PathBuf>,
     pub arguments: ArgMatches<'a>,
+    pub repository: Option<Repository>,
 }
 
 impl<'a> Context<'a> {
@@ -34,10 +36,13 @@ impl<'a> Context<'a> {
             .map(|entry| entry.path())
             .collect::<Vec<PathBuf>>();
 
+        let repository: Option<Repository> = Repository::discover(&current_dir).ok();
+
         Context {
-            current_dir,
             arguments,
+            current_dir,
             dir_files,
+            repository,
         }
     }
 
diff --git a/src/modules/directory.rs b/src/modules/directory.rs
index 83e5613bd..d5fbc88c9 100644
--- a/src/modules/directory.rs
+++ b/src/modules/directory.rs
@@ -1,5 +1,4 @@
 use ansi_term::Color;
-use git2::Repository;
 use std::path::Path;
 
 use super::Segment;
@@ -23,17 +22,17 @@ pub fn segment(context: &Context) -> Option<Segment> {
     let current_dir = &context.current_dir;
 
     let dir_string;
-    if let Ok(repo) = git2::Repository::discover(current_dir) {
+    if let Some(repo) = &context.repository {
         // Contract the path to the git repo root
-        let repo_root = get_repo_root(&repo);
+        let repo_root = repo.workdir().unwrap();
         let repo_folder_name = repo_root.file_name().unwrap().to_str().unwrap();
 
-        dir_string = contract_path(current_dir, repo_root, repo_folder_name);
+        dir_string = contract_path(&current_dir, repo_root, repo_folder_name);
     } else {
         // Contract the path to the home directory
         let home_dir = dirs::home_dir().unwrap();
 
-        dir_string = contract_path(current_dir, &home_dir, HOME_SYMBOL);
+        dir_string = contract_path(&current_dir, &home_dir, HOME_SYMBOL);
     }
 
     // Truncate the dir string to the maximum number of path components
@@ -46,17 +45,6 @@ pub fn segment(context: &Context) -> Option<Segment> {
     Some(segment)
 }
 
-/// Get the root directory of a git repo
-fn get_repo_root(repo: &Repository) -> &Path {
-    if repo.is_bare() {
-        // Bare repos will return the repo root
-        repo.path()
-    } else {
-        // Non-bare repos will return the path of `.git`
-        repo.path().parent().unwrap()
-    }
-}
-
 /// Contract the root component of a path
 ///
 /// Replaces the `top_level_path` in a given `full_path` with the provided
diff --git a/src/modules/git_branch.rs b/src/modules/git_branch.rs
new file mode 100644
index 000000000..586724b2e
--- /dev/null
+++ b/src/modules/git_branch.rs
@@ -0,0 +1,46 @@
+use ansi_term::Color;
+use git2::Repository;
+
+use super::Segment;
+use crate::context::Context;
+
+/// 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> {
+    if context.repository.is_none() {
+        return None;
+    }
+
+    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;
+
+            // 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 segment = Segment::new("git_branch");
+            segment
+                .set_prefix(Some(Box::new(segment_prefix)))
+                .set_style(SEGMENT_COLOR.bold())
+                .set_value(branch_name);
+
+            Some(segment)
+        }
+        Err(_e) => None,
+    }
+}
+
+fn get_current_branch(repository: &Repository) -> Result<String, git2::Error> {
+    let head = repository.head()?;
+    let head_name = head.shorthand();
+    match head_name {
+        Some(name) => Ok(name.to_string()),
+        None => Err(git2::Error::from_str("No branch name found")),
+    }
+}
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index 8e034331b..4ed912adc 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -1,5 +1,6 @@
 mod character;
 mod directory;
+mod git_branch;
 mod line_break;
 mod nodejs;
 mod python;
@@ -16,6 +17,7 @@ pub fn handle(module: &str, context: &Context) -> Option<Segment> {
         "rust" | "rustlang" => rust::segment(context),
         "python" => python::segment(context),
         "line_break" => line_break::segment(context),
+        "git_branch" => git_branch::segment(context),
 
         _ => panic!("Unknown module: {}", module),
     }
diff --git a/src/modules/nodejs.rs b/src/modules/nodejs.rs
index d8431038a..818a0cf5c 100644
--- a/src/modules/nodejs.rs
+++ b/src/modules/nodejs.rs
@@ -23,7 +23,7 @@ pub fn segment(context: &Context) -> Option<Segment> {
             const SEGMENT_COLOR: Color = Color::Green;
 
             let mut segment = Segment::new("node");
-            segment.set_style(SEGMENT_COLOR);
+            segment.set_style(SEGMENT_COLOR.bold());
 
             let formatted_version = node_version.trim();
             segment.set_value(format!("{} {}", NODE_CHAR, formatted_version));
diff --git a/src/modules/python.rs b/src/modules/python.rs
index 3d182bb28..afd62690a 100644
--- a/src/modules/python.rs
+++ b/src/modules/python.rs
@@ -23,7 +23,7 @@ pub fn segment(context: &Context) -> Option<Segment> {
             const SEGMENT_COLOR: Color = Color::Yellow;
 
             let mut segment = Segment::new("python");
-            segment.set_style(SEGMENT_COLOR);
+            segment.set_style(SEGMENT_COLOR.bold());
 
             let formatted_version = format_python_version(python_version);
             segment.set_value(format!("{} {}", PYTHON_CHAR, formatted_version));
diff --git a/src/modules/rust.rs b/src/modules/rust.rs
index 28a4a83f5..14af3b854 100644
--- a/src/modules/rust.rs
+++ b/src/modules/rust.rs
@@ -21,7 +21,7 @@ pub fn segment(context: &Context) -> Option<Segment> {
             const SEGMENT_COLOR: Color = Color::Red;
 
             let mut segment = Segment::new("rust");
-            segment.set_style(SEGMENT_COLOR);
+            segment.set_style(SEGMENT_COLOR.bold());
 
             let formatted_version = format_rustc_version(rust_version);
             segment.set_value(format!("{} {}", RUST_CHAR, formatted_version));
diff --git a/src/print.rs b/src/print.rs
index 836795d42..b175252eb 100644
--- a/src/print.rs
+++ b/src/print.rs
@@ -7,6 +7,7 @@ use crate::modules;
 pub fn prompt(args: ArgMatches) {
     let prompt_order = vec![
         "directory",
+        "git_branch",
         "nodejs",
         "rust",
         "python",