Implement the git status module (#45)

7 files changed, 168 insertions, 27 deletions

diff --git a/src/context.rs b/src/context.rs
index 63ca28e76..fbf85268c 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -10,6 +10,7 @@ pub struct Context<'a> {
     pub dir_files: Vec<PathBuf>,
     pub arguments: ArgMatches<'a>,
     pub repo_root: Option<PathBuf>,
+    pub branch_name: Option<String>,
 }
 
 impl<'a> Context<'a> {
@@ -37,15 +38,20 @@ impl<'a> Context<'a> {
             .map(|entry| entry.path())
             .collect::<Vec<PathBuf>>();
 
-        let repo_root: Option<PathBuf> = Repository::discover(&current_dir)
-            .ok()
+        let repository = Repository::discover(&current_dir).ok();
+        let repo_root = repository
+            .as_ref()
             .and_then(|repo| repo.workdir().map(|repo| repo.to_path_buf()));
+        let branch_name = repository
+            .as_ref()
+            .and_then(|repo| get_current_branch(&repo));
 
         Context {
             arguments,
             current_dir,
             dir_files,
             repo_root,
+            branch_name,
         }
     }
 
@@ -192,3 +198,10 @@ mod tests {
         assert_eq!(passing_criteria.scan(), true);
     }
 }
+
+fn get_current_branch(repository: &Repository) -> Option<String> {
+    let head = repository.head().ok()?;
+    let shorthand = head.shorthand();
+
+    shorthand.map(|branch| branch.to_string())
+}
diff --git a/src/main.rs b/src/main.rs
index d5422a129..0d873868a 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,6 +4,7 @@ extern crate clap;
 extern crate ansi_term;
 extern crate dirs;
 extern crate git2;
+extern crate pretty_env_logger;
 
 mod context;
 mod module;
@@ -14,6 +15,8 @@ mod segment;
 use clap::{App, Arg};
 
 fn main() {
+    pretty_env_logger::init();
+
     let args = App::new("Starship")
         .about("The cross-shell prompt for astronauts. ✨🚀")
         // pull the version number from Cargo.toml
diff --git a/src/module.rs b/src/module.rs
index 903d5cf54..43ff78c1e 100644
--- a/src/module.rs
+++ b/src/module.rs
@@ -48,6 +48,11 @@ impl Module {
         self.segments.last_mut().unwrap()
     }
 
+    /// Whether a module has any segments
+    pub fn is_empty(&self) -> bool {
+        self.segments.is_empty()
+    }
+
     /// Get the module's prefix
     pub fn get_prefix(&mut self) -> &mut ModuleAffix {
         &mut self.prefix
diff --git a/src/modules/git_branch.rs b/src/modules/git_branch.rs
index 22b934da3..500cc10b8 100644
--- a/src/modules/git_branch.rs
+++ b/src/modules/git_branch.rs
@@ -1,5 +1,4 @@
 use ansi_term::Color;
-use git2::Repository;
 
 use super::{Context, Module};
 
@@ -7,32 +6,17 @@ use super::{Context, Module};
 ///
 /// Will display the branch name if the current directory is a git repo
 pub fn segment(context: &Context) -> Option<Module> {
-    let repo_root = context.repo_root.as_ref()?;
-    let repository = Repository::open(repo_root).ok()?;
+    let branch_name = context.branch_name.as_ref()?;
 
-    match get_current_branch(&repository) {
-        Ok(branch_name) => {
-            const GIT_BRANCH_CHAR: &str = " ";
-            let segment_color = Color::Purple.bold();
+    const GIT_BRANCH_CHAR: &str = " ";
+    let segment_color = Color::Purple.bold();
 
-            let mut module = Module::new("git_branch");
-            module.set_style(segment_color);
-            module.get_prefix().set_value("in ");
+    let mut module = Module::new("git_branch");
+    module.set_style(segment_color);
+    module.get_prefix().set_value("on ");
 
-            module.new_segment("branch_char", GIT_BRANCH_CHAR);
-            module.new_segment("branch_name", branch_name);
+    module.new_segment("branch_char", GIT_BRANCH_CHAR);
+    module.new_segment("branch_name", branch_name.to_string());
 
-            Some(module)
-        }
-        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")),
-    }
+    Some(module)
 }
diff --git a/src/modules/git_status.rs b/src/modules/git_status.rs
new file mode 100644
index 000000000..453edd9a5
--- /dev/null
+++ b/src/modules/git_status.rs
@@ -0,0 +1,133 @@
+use ansi_term::Color;
+use git2::{Repository, Status};
+
+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
+/// By default, the following symbols will be used to represent the repo's status:
+///   - `=` – This branch has merge conflicts
+///   - `⇡` – This branch is ahead of the branch being tracked
+///   - `⇡` – This branch is behind of the branch being tracked
+///   - `⇕` – This branch has diverged from the branch being tracked
+///   - `?` — There are untracked files in the working directory
+///   - `$` — A stash exists for the repository
+///   - `!` — There are file modifications in the working directory
+///   - `+` — A new file has been added to the staging area
+///   - `»` — A renamed file has been added to the staging area
+///   - `✘` — A file's deletion has been added to the staging area
+pub fn segment(context: &Context) -> Option<Module> {
+    // This is the order that the sections will appear in
+    const GIT_STATUS_CONFLICTED: &str = "=";
+    const GIT_STATUS_AHEAD: &str = "⇡";
+    const GIT_STATUS_BEHIND: &str = "⇣";
+    const GIT_STATUS_DIVERGED: &str = "⇕";
+    const GIT_STATUS_UNTRACKED: &str = "?";
+    const GIT_STATUS_STASHED: &str = "$";
+    const GIT_STATUS_MODIFIED: &str = "!";
+    const GIT_STATUS_ADDED: &str = "+";
+    const GIT_STATUS_RENAMED: &str = "»";
+    const GIT_STATUS_DELETED: &str = "✘";
+
+    let branch_name = context.branch_name.as_ref()?;
+    let repo_root = context.repo_root.as_ref()?;
+    let repository = Repository::open(repo_root).ok()?;
+
+    let module_style = Color::Red.bold();
+    let mut module = Module::new("git_status");
+    module.get_prefix().set_value("[").set_style(module_style);
+    module.get_suffix().set_value("] ").set_style(module_style);
+    module.set_style(module_style);
+
+    let ahead_behind = get_ahead_behind(&repository, &branch_name);
+    log::debug!("Repo ahead/behind: {:?}", ahead_behind);
+    let stash_object = repository.revparse_single("refs/stash");
+    log::debug!("Stash object: {:?}", stash_object);
+    let repo_status = get_repo_status(&repository);
+    log::debug!("Repo status: {:?}", repo_status);
+
+    // Add the conflicted segment
+    if let Ok(repo_status) = repo_status {
+        if repo_status.is_conflicted() {
+            module.new_segment("conflicted", GIT_STATUS_CONFLICTED);
+        }
+    }
+
+    // Add the ahead/behind segment
+    if let Ok((ahead, behind)) = ahead_behind {
+        if ahead > 0 && behind > 0 {
+            module.new_segment("diverged", GIT_STATUS_DIVERGED);
+        } else if ahead > 0 {
+            module.new_segment("ahead", GIT_STATUS_AHEAD);
+        } else if behind > 0 {
+            module.new_segment("behind", GIT_STATUS_BEHIND);
+        }
+    }
+
+    // Add the stashed segment
+    if stash_object.is_ok() {
+        module.new_segment("stashed", GIT_STATUS_STASHED);
+    }
+
+    // Add all remaining status segments
+    if let Ok(repo_status) = repo_status {
+        if repo_status.is_wt_deleted() || repo_status.is_index_deleted() {
+            module.new_segment("deleted", GIT_STATUS_DELETED);
+        }
+
+        if repo_status.is_wt_renamed() || repo_status.is_index_renamed() {
+            module.new_segment("renamed", GIT_STATUS_RENAMED);
+        }
+
+        if repo_status.is_wt_modified() {
+            module.new_segment("modified", GIT_STATUS_MODIFIED);
+        }
+
+        if repo_status.is_index_modified() || repo_status.is_index_new() {
+            module.new_segment("staged", GIT_STATUS_ADDED);
+        }
+
+        if repo_status.is_wt_new() {
+            module.new_segment("untracked", GIT_STATUS_UNTRACKED);
+        }
+    }
+
+    if module.is_empty() {
+        None
+    } else {
+        Some(module)
+    }
+}
+
+/// Gets the bitflags associated with the repo's git status
+fn get_repo_status(repository: &Repository) -> Result<Status, git2::Error> {
+    let mut status_options = git2::StatusOptions::new();
+    status_options.include_untracked(true);
+
+    let repo_file_statuses = repository.statuses(Some(&mut status_options))?;
+
+    // Statuses are stored as bitflags, so use BitOr to join them all into a single value
+    let repo_status: Status = repo_file_statuses.iter().map(|e| e.status()).collect();
+    if repo_status.is_empty() {
+        return Err(git2::Error::from_str("Repo has no status"));
+    }
+
+    Ok(repo_status)
+}
+
+/// Compares the current branch with the branch it is tracking to determine how
+/// far ahead or behind it is in relation
+fn get_ahead_behind(
+    repository: &Repository,
+    branch_name: &str,
+) -> Result<(usize, usize), git2::Error> {
+    let branch_object = repository.revparse_single(branch_name)?;
+    let tracking_branch_name = format!("{}@{{upstream}}", branch_name);
+    let tracking_object = repository.revparse_single(&tracking_branch_name)?;
+
+    let branch_oid = branch_object.id();
+    let tracking_oid = tracking_object.id();
+
+    repository.graph_ahead_behind(branch_oid, tracking_oid)
+}
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index 2a9f433bb..5b3e7167f 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -1,6 +1,7 @@
 mod character;
 mod directory;
 mod git_branch;
+mod git_status;
 mod go;
 mod line_break;
 mod nodejs;
@@ -22,6 +23,7 @@ pub fn handle(module: &str, context: &Context) -> Option<Module> {
         "line_break" => line_break::segment(context),
         "package" => package::segment(context),
         "git_branch" => git_branch::segment(context),
+        "git_status" => git_status::segment(context),
 
         _ => panic!("Unknown module: {}", module),
     }
diff --git a/src/print.rs b/src/print.rs
index 2f26ddcc1..f4e5ae950 100644
--- a/src/print.rs
+++ b/src/print.rs
@@ -10,6 +10,7 @@ pub fn prompt(args: ArgMatches) {
     let prompt_order = vec![
         "directory",
         "git_branch",
+        "git_status",
         "package",
         "nodejs",
         "rust",