feat: added truncation_length/symbol to git_branch (#268)

Git branches can become very long (e.g. gitlab auto-generated branch
names), thus it would be nice to be able to truncate them to keep your
prompt lenght in line.

This patch adds two new options to the git_branch module:
* truncation_length: The amount of graphemes to of a gitbranch to
truncate to
* truncation_symbol: The symbol that should be used to indicate that a
branch name was trunctated

To be able to correctly work with UTF-8 graphemes, unicode-segmentation
was added as a dependency.

6 files changed, 227 insertions, 6 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 53dec7a0f..715bb0948 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -754,6 +754,7 @@ dependencies = [
  "serde_json 1.0.40 (registry+https://github.com/rust-lang/crates.io-index)",
  "tempfile 3.1.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "toml 0.5.3 (registry+https://github.com/rust-lang/crates.io-index)",
+ "unicode-segmentation 1.3.0 (registry+https://github.com/rust-lang/crates.io-index)",
 ]
 
 [[package]]
@@ -875,6 +876,11 @@ dependencies = [
 ]
 
 [[package]]
+name = "unicode-segmentation"
+version = "1.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+
+[[package]]
 name = "unicode-width"
 version = "0.1.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1065,6 +1071,7 @@ dependencies = [
 "checksum ucd-util 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)" = "535c204ee4d8434478593480b8f86ab45ec9aae0e83c568ca81abf0fd0e88f86"
 "checksum unicode-bidi 0.3.4 (registry+https://github.com/rust-lang/crates.io-index)" = "49f2bd0c6468a8230e1db229cff8029217cf623c767ea5d60bfbd42729ea54d5"
 "checksum unicode-normalization 0.1.8 (registry+https://github.com/rust-lang/crates.io-index)" = "141339a08b982d942be2ca06ff8b076563cbe223d1befd5450716790d44e2426"
+"checksum unicode-segmentation 1.3.0 (registry+https://github.com/rust-lang/crates.io-index)" = "1967f4cdfc355b37fd76d2a954fb2ed3871034eb4f26d60537d88795cfc332a9"
 "checksum unicode-width 0.1.5 (registry+https://github.com/rust-lang/crates.io-index)" = "882386231c45df4700b275c7ff55b6f3698780a650026380e72dabe76fa46526"
 "checksum unicode-xid 0.1.0 (registry+https://github.com/rust-lang/crates.io-index)" = "fc72304796d0818e357ead4e000d19c9c174ab23dc11093ac919054d20a6a7fc"
 "checksum uom 0.23.0 (registry+https://github.com/rust-lang/crates.io-index)" = "347fe3ff20637a62ab9749a5c90d167302bcbdab77ec961dda7f62a5ca6d368a"
diff --git a/Cargo.toml b/Cargo.toml
index 9c3d08bdd..3249c8ce3 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -39,6 +39,7 @@ log = "0.4.8"
 battery = { version = "0.7.4", optional = true }
 lazy_static = "1.4.0"
 path-slash = "0.1.1"
+unicode-segmentation = "1.3.0"
 
 [dev-dependencies]
 tempfile = "3.1.0"
diff --git a/docs/config/README.md b/docs/config/README.md
index a91738d4a..732e1528d 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -211,10 +211,12 @@ The `git_branch` module shows the active branch of the repo in your current dire
 
 ### Options
 
-| Variable   | Default | Description                                                                   |
-| ---------- | ------- | ----------------------------------------------------------------------------- |
-| `symbol`   | `" "`  | The symbol used before the branch name of the repo in your current directory. |
-| `disabled` | `false` | Disables the `git_branch` module.                                             |
+| Variable            | Default    | Description                                                                           |
+| ------------------- | ---------- | ------------------------------------------------------------------------------------- |
+| `symbol`            | `" "`     | The symbol used before the branch name of the repo in your current directory.         |
+| `disabled`          | `false`    | Disables the `git_branch` module.                                                     |
+| `truncation_length` | `2^63 - 1` | Truncates a git branch to X graphemes                                                 |
+| `truncation_symbol` | `"…"`      | The symbol used to indicate a branch name was truncated. You can use "" for no symbol |
 
 ### Example
 
@@ -223,6 +225,8 @@ The `git_branch` module shows the active branch of the repo in your current dire
 
 [git_branch]
 symbol = "🌱 "
+truncation_length = "4"
+truncation_symbol = ""
 ```
 
 ## Git Status
diff --git a/src/modules/git_branch.rs b/src/modules/git_branch.rs
index e3a12e272..120c58d4e 100644
--- a/src/modules/git_branch.rs
+++ b/src/modules/git_branch.rs
@@ -1,4 +1,5 @@
 use ansi_term::Color;
+use unicode_segmentation::UnicodeSegmentation;
 
 use super::{Context, Module};
 
@@ -8,15 +9,54 @@ use super::{Context, Module};
 pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     const GIT_BRANCH_CHAR: &str = " ";
 
-    let branch_name = context.branch_name.as_ref()?;
     let segment_color = Color::Purple.bold();
 
     let mut module = context.new_module("git_branch")?;
     module.set_style(segment_color);
     module.get_prefix().set_value("on ");
 
+    let unsafe_truncation_length = module
+        .config_value_i64("truncation_length")
+        .unwrap_or(std::i64::MAX);
+    let truncation_symbol = get_graphemes(
+        module.config_value_str("truncation_symbol").unwrap_or("…"),
+        1,
+    );
+
     module.new_segment("symbol", GIT_BRANCH_CHAR);
-    module.new_segment("name", branch_name);
+
+    // TODO: Once error handling is implemented, warn the user if their config
+    // truncation length is nonsensical
+    let len = if unsafe_truncation_length <= 0 {
+        log::debug!(
+            "[WARN]: \"truncation_length\" should be a positive value, found {}",
+            unsafe_truncation_length
+        );
+        std::usize::MAX
+    } else {
+        unsafe_truncation_length as usize
+    };
+    let branch_name = context.branch_name.as_ref()?;
+    let truncated_graphemes = get_graphemes(&branch_name, len);
+    // The truncation symbol should only be added if we truncated
+    let truncated_and_symbol = if len < graphemes_len(&branch_name) {
+        truncated_graphemes + &truncation_symbol
+    } else {
+        truncated_graphemes
+    };
+
+    module.new_segment("name", &truncated_and_symbol);
 
     Some(module)
 }
+
+fn get_graphemes(text: &str, length: usize) -> String {
+    UnicodeSegmentation::graphemes(text, true)
+        .take(length)
+        .collect::<Vec<&str>>()
+        .concat()
+}
+
+fn graphemes_len(text: &str) -> usize {
+    UnicodeSegmentation::graphemes(&text[..], true).count()
+}
diff --git a/tests/testsuite/git_branch.rs b/tests/testsuite/git_branch.rs
new file mode 100644
index 000000000..16d889b4d
--- /dev/null
+++ b/tests/testsuite/git_branch.rs
@@ -0,0 +1,168 @@
+use ansi_term::Color;
+use git2::Repository;
+use std::env;
+use std::io;
+use std::process::Command;
+
+use crate::common::{self, TestCommand};
+
+#[test]
+fn test_changed_truncation_symbol() -> io::Result<()> {
+    test_truncate_length_with_config(
+        "1337_hello_world",
+        15,
+        "1337_hello_worl",
+        "%",
+        "truncation_symbol = \"%\"",
+    )
+}
+
+#[test]
+fn test_no_truncation_symbol() -> io::Result<()> {
+    test_truncate_length_with_config(
+        "1337_hello_world",
+        15,
+        "1337_hello_worl",
+        "",
+        "truncation_symbol = \"\"",
+    )
+}
+
+#[test]
+fn test_multi_char_truncation_symbol() -> io::Result<()> {
+    test_truncate_length_with_config(
+        "1337_hello_world",
+        15,
+        "1337_hello_worl",
+        "a",
+        "truncation_symbol = \"apple\"",
+    )
+}
+
+#[test]
+fn test_ascii_boundary_below() -> io::Result<()> {
+    test_truncate_length("1337_hello_world", 15, "1337_hello_worl", "…")
+}
+
+#[test]
+fn test_ascii_boundary_on() -> io::Result<()> {
+    test_truncate_length("1337_hello_world", 16, "1337_hello_world", "")
+}
+
+#[test]
+fn test_ascii_boundary_above() -> io::Result<()> {
+    test_truncate_length("1337_hello_world", 17, "1337_hello_world", "")
+}
+
+#[test]
+fn test_one() -> io::Result<()> {
+    test_truncate_length("1337_hello_world", 1, "1", "…")
+}
+
+#[test]
+fn test_zero() -> io::Result<()> {
+    test_truncate_length("1337_hello_world", 0, "1337_hello_world", "")
+}
+
+#[test]
+fn test_negative() -> io::Result<()> {
+    test_truncate_length("1337_hello_world", -1, "1337_hello_world", "")
+}
+
+#[test]
+fn test_hindi_truncation() -> io::Result<()> {
+    test_truncate_length("नमस्ते", 3, "नमस्", "…")
+}
+
+#[test]
+fn test_hindi_truncation2() -> io::Result<()> {
+    test_truncate_length("नमस्त", 3, "नमस्", "…")
+}
+
+#[test]
+fn test_japanese_truncation() -> io::Result<()> {
+    test_truncate_length("がんばってね", 4, "がんばっ", "…")
+}
+
+fn test_truncate_length(
+    branch_name: &str,
+    truncate_length: i64,
+    expected_name: &str,
+    truncation_symbol: &str,
+) -> io::Result<()> {
+    return test_truncate_length_with_config(
+        branch_name,
+        truncate_length,
+        expected_name,
+        truncation_symbol,
+        "",
+    );
+}
+
+fn test_truncate_length_with_config(
+    branch_name: &str,
+    truncate_length: i64,
+    expected_name: &str,
+    truncation_symbol: &str,
+    config_options: &str,
+) -> io::Result<()> {
+    let fixture_repo_dir = 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(&["checkout", "-b", branch_name])
+        .current_dir(repo_dir.as_path())
+        .output()?;
+
+    let output = common::render_module("git_branch")
+        .use_config(
+            toml::from_str(&format!(
+                "
+                    [git_branch]
+                        truncation_length = {}
+                        {}
+                ",
+                truncate_length, config_options
+            ))
+            .unwrap(),
+        )
+        .arg("--path")
+        .arg(repo_dir)
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+
+    let expected = format!(
+        "on {} ",
+        Color::Purple
+            .bold()
+            .paint(format!("\u{e0a0} {}{}", expected_name, truncation_symbol)),
+    );
+    assert_eq!(expected, actual);
+    Ok(())
+}
+
+fn create_fixture_repo() -> io::Result<std::path::PathBuf> {
+    let fixture_repo_dir = common::new_tempdir()?.path().join("fixture");
+    let fixture = env::current_dir()?.join("tests/fixtures/rocket.bundle");
+
+    Command::new("git")
+        .args(&["config", "--global", "user.email", "starship@example.com"])
+        .output()?;
+
+    Command::new("git")
+        .args(&["config", "--global", "user.name", "starship"])
+        .output()?;
+
+    Command::new("git")
+        .args(&[
+            "clone",
+            "-b",
+            "master",
+            &fixture.to_str().unwrap(),
+            fixture_repo_dir.to_str().unwrap(),
+        ])
+        .output()?;
+
+    Ok(fixture_repo_dir)
+}
diff --git a/tests/testsuite/main.rs b/tests/testsuite/main.rs
index ed0289835..cd96b8e85 100644
--- a/tests/testsuite/main.rs
+++ b/tests/testsuite/main.rs
@@ -3,6 +3,7 @@ mod cmd_duration;
 mod common;
 mod configuration;
 mod directory;
+mod git_branch;
 mod git_status;
 mod golang;
 mod jobs;