feat: Add config for ahead/behind count of tracked branch (#281)

Add a configuration option (show_sync_count) to the git_status module that will show/hide the counts ahead/behind of the tracked branch. Currently have this default to false, and would opt-in to show this information.

3 files changed, 144 insertions, 24 deletions

diff --git a/docs/config/README.md b/docs/config/README.md
index e7e8e5f4b..944545246 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -237,19 +237,20 @@ current directory.
 
 ### Options
 
-| Variable     | Default | Description                                             |
-| ------------ | ------- | ------------------------------------------------------- |
-| `conflicted` | `"="`   | This branch has merge conflicts.                        |
-| `ahead`      | `"⇡"`   | This branch is ahead of the branch being tracked.       |
-| `behind`     | `"⇣"`   | This branch is behind of the branch being tracked.      |
-| `diverged`   | `"⇕"`   | This branch has diverged from the branch being tracked. |
-| `untracked`  | `"?"`   | There are untracked files in the working directory.     |
-| `stashed`    | `"$"`   | A stash exists for the local repository.                |
-| `modified`   | `"!"`   | There are file modifications in the working directory.  |
-| `staged`     | `"+"`   | A new file has been added to the staging area.          |
-| `renamed`    | `"»"`   | A renamed file has been added to the staging area.      |
-| `deleted`    | `"✘"`   | A file's deletion has been added to the staging area.   |
-| `disabled`   | `false` | Disables the `git_status` module.                       |
+| Variable          | Default | Description                                             |
+| ----------------- | ------- | ------------------------------------------------------- |
+| `conflicted`      | `"="`   | This branch has merge conflicts.                        |
+| `ahead`           | `"⇡"`   | This branch is ahead of the branch being tracked.       |
+| `behind`          | `"⇣"`   | This branch is behind of the branch being tracked.      |
+| `diverged`        | `"⇕"`   | This branch has diverged from the branch being tracked. |
+| `untracked`       | `"?"`   | There are untracked files in the working directory.     |
+| `stashed`         | `"$"`   | A stash exists for the local repository.                |
+| `modified`        | `"!"`   | There are file modifications in the working directory.  |
+| `staged`          | `"+"`   | A new file has been added to the staging area.          |
+| `renamed`         | `"»"`   | A renamed file has been added to the staging area.      |
+| `deleted`         | `"✘"`   | A file's deletion has been added to the staging area.   |
+| `show_sync_count` | `false` | Show ahead/behind count of the branch being tracked.    |
+| `disabled`        | `false` | Disables the `git_status` module.                       |
 
 ### Example
 
diff --git a/src/modules/git_status.rs b/src/modules/git_status.rs
index 2a637957e..646f01f07 100644
--- a/src/modules/git_status.rs
+++ b/src/modules/git_status.rs
@@ -36,6 +36,9 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
 
     let module_style = Color::Red.bold();
     let mut module = context.new_module("git_status")?;
+
+    let show_sync_count = module.config_value_bool("show_sync_count").unwrap_or(false);
+
     module.get_prefix().set_value("[").set_style(module_style);
     module.get_suffix().set_value("] ").set_style(module_style);
     module.set_style(module_style);
@@ -66,17 +69,37 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
 
     // Add the ahead/behind segment
     if let Ok((ahead, behind)) = ahead_behind {
-        let ahead_segment = format!("{}{}", GIT_STATUS_AHEAD, ahead);
-        let behind_segment = format!("{}{}", GIT_STATUS_BEHIND, behind);
+        let add_ahead = |m: &mut Module<'a>| {
+            m.new_segment("ahead", GIT_STATUS_AHEAD);
+
+            if show_sync_count {
+                m.new_segment("ahead_count", &ahead.to_string());
+            }
+        };
+
+        let add_behind = |m: &mut Module<'a>| {
+            m.new_segment("behind", GIT_STATUS_BEHIND);
+
+            if show_sync_count {
+                m.new_segment("behind_count", &behind.to_string());
+            }
+        };
 
         if ahead > 0 && behind > 0 {
             module.new_segment("diverged", GIT_STATUS_DIVERGED);
-            module.new_segment("ahead", ahead_segment.as_str());
-            module.new_segment("behind", behind_segment.as_str());
-        } else if ahead > 0 {
-            module.new_segment("ahead", ahead_segment.as_str());
-        } else if behind > 0 {
-            module.new_segment("behind", behind_segment.as_str());
+
+            if show_sync_count {
+                add_ahead(&mut module);
+                add_behind(&mut module);
+            }
+        }
+
+        if ahead > 0 && behind == 0 {
+            add_ahead(&mut module);
+        }
+
+        if behind > 0 && ahead == 0 {
+            add_behind(&mut module);
         }
     }
 
diff --git a/tests/testsuite/git_status.rs b/tests/testsuite/git_status.rs
index 00866247b..4b75d0a05 100644
--- a/tests/testsuite/git_status.rs
+++ b/tests/testsuite/git_status.rs
@@ -5,11 +5,11 @@ use std::fs::{self, File};
 use std::io;
 use std::process::Command;
 
-use crate::common;
+use crate::common::{self, TestCommand};
 
 #[test]
 #[ignore]
-fn shows_behind_count() -> io::Result<()> {
+fn shows_behind() -> io::Result<()> {
     let fixture_repo_dir = common::create_fixture_repo()?;
     let repo_dir = common::new_tempdir()?.path().join("rocket");
 
@@ -25,6 +25,35 @@ fn shows_behind_count() -> io::Result<()> {
         .arg(repo_dir)
         .output()?;
     let actual = String::from_utf8(output.stdout).unwrap();
+    let expected = Color::Red.bold().paint(format!("[{}] ", "⇣")).to_string();
+
+    assert_eq!(expected, actual);
+
+    Ok(())
+}
+
+#[test]
+#[ignore]
+fn shows_behind_with_count() -> io::Result<()> {
+    let fixture_repo_dir = common::create_fixture_repo()?;
+    let repo_dir = common::new_tempdir()?.path().join("rocket");
+
+    Repository::clone(fixture_repo_dir.to_str().unwrap(), &repo_dir.as_path()).unwrap();
+
+    Command::new("git")
+        .args(&["reset", "--hard", "HEAD^"])
+        .current_dir(repo_dir.as_path())
+        .output()?;
+
+    let output = common::render_module("git_status")
+        .use_config(toml::toml! {
+            [git_status]
+            show_sync_count = true
+        })
+        .arg("--path")
+        .arg(repo_dir)
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
     let expected = Color::Red
         .bold()
         .paint(format!("[{}] ", "⇣1"))
@@ -37,7 +66,34 @@ fn shows_behind_count() -> io::Result<()> {
 
 #[test]
 #[ignore]
-fn shows_ahead_count() -> io::Result<()> {
+fn shows_ahead() -> io::Result<()> {
+    let fixture_repo_dir = common::create_fixture_repo()?;
+    let repo_dir = common::new_tempdir()?.path().join("rocket");
+
+    Repository::clone(fixture_repo_dir.to_str().unwrap(), &repo_dir.as_path()).unwrap();
+
+    File::create(repo_dir.join("readme.md"))?;
+
+    Command::new("git")
+        .args(&["commit", "-am", "Update readme"])
+        .current_dir(&repo_dir)
+        .output()?;
+
+    let output = common::render_module("git_status")
+        .arg("--path")
+        .arg(repo_dir)
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+    let expected = Color::Red.bold().paint(format!("[{}] ", "⇡")).to_string();
+
+    assert_eq!(expected, actual);
+
+    Ok(())
+}
+
+#[test]
+#[ignore]
+fn shows_ahead_with_count() -> io::Result<()> {
     let fixture_repo_dir = common::create_fixture_repo()?;
     let repo_dir = common::new_tempdir()?.path().join("rocket");
 
@@ -51,6 +107,10 @@ fn shows_ahead_count() -> io::Result<()> {
         .output()?;
 
     let output = common::render_module("git_status")
+        .use_config(toml::toml! {
+            [git_status]
+            show_sync_count = true
+        })
         .arg("--path")
         .arg(repo_dir)
         .output()?;
@@ -90,6 +150,42 @@ fn shows_diverged() -> io::Result<()> {
         .arg(repo_dir)
         .output()?;
     let actual = String::from_utf8(output.stdout).unwrap();
+    let expected = Color::Red.bold().paint(format!("[{}] ", "⇕")).to_string();
+
+    assert_eq!(expected, actual);
+
+    Ok(())
+}
+
+#[test]
+#[ignore]
+fn shows_diverged_with_count() -> io::Result<()> {
+    let fixture_repo_dir = common::create_fixture_repo()?;
+    let repo_dir = common::new_tempdir()?.path().join("rocket");
+
+    Repository::clone(fixture_repo_dir.to_str().unwrap(), &repo_dir.as_path()).unwrap();
+
+    Command::new("git")
+        .args(&["reset", "--hard", "HEAD^"])
+        .current_dir(repo_dir.as_path())
+        .output()?;
+
+    fs::write(repo_dir.join("Cargo.toml"), " ")?;
+
+    Command::new("git")
+        .args(&["commit", "-am", "Update readme"])
+        .current_dir(repo_dir.as_path())
+        .output()?;
+
+    let output = common::render_module("git_status")
+        .use_config(toml::toml! {
+            [git_status]
+            show_sync_count = true
+        })
+        .arg("--path")
+        .arg(repo_dir)
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
     let expected = Color::Red
         .bold()
         .paint(format!("[{}] ", "⇕⇡1⇣1"))