feat: Add the hg_branch module (#569)

12 files changed, 383 insertions, 0 deletions

diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml
index aa96610a0..7308948cf 100644
--- a/.github/workflows/workflow.yml
+++ b/.github/workflows/workflow.yml
@@ -124,6 +124,17 @@ jobs:
         with:
           dotnet-version: "2.2.402"
 
+      # Install Mercurial (pre-installed on linux, installed from pip on macos
+      # and from choco on windows),
+      - name: Install Mercurial (macos)
+        if: matrix.os == 'macOS-latest'
+        env:
+          HGPYTHON3: 1
+        run: pip install mercurial
+      - name: Install Mercurial (windows)
+        if: matrix.os == 'windows-latest'
+        run: choco install hg
+
       # Run the ignored tests that expect the above setup
       - name: Run all tests
         uses: actions-rs/cargo@v1
diff --git a/docs/config/README.md b/docs/config/README.md
index 797c438f5..cfc94e484 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -93,6 +93,7 @@ prompt_order = [
     "git_branch",
     "git_state",
     "git_status",
+    "hg_branch",
     "package",
     "dotnet",
     "golang",
@@ -559,6 +560,31 @@ The module will be shown if any of the following conditions are met:
 symbol = "🏎💨 "
 ```
 
+## Mercurial Branch
+
+The `hg_branch` module shows the active branch of the repo in your current directory.
+
+### Options
+
+| Variable            | Default         | Description                                                                                  |
+| ------------------- | --------------- | -------------------------------------------------------------------------------------------- |
+| `symbol`            | `" "`          | The symbol used before the hg bookmark or branch name of the repo in your current directory. |
+| `truncation_length` | `2^63 - 1`      | Truncates the hg branch name to X graphemes                                                  |
+| `truncation_symbol` | `"…"`           | The symbol used to indicate a branch name was truncated.                                     |
+| `style`             | `"bold purple"` | The style for the module.                                                                    |
+| `disabled`          | `true`          | Disables the `hg_branch` module.                                                             |
+
+### Example
+
+```toml
+# ~/.config/starship.toml
+
+[hg_branch]
+symbol = "🌱 "
+truncation_length = 4
+truncation_symbol = ""
+```
+
 ## Hostname
 
 The `hostname` module shows the system hostname.
diff --git a/src/configs/hg_branch.rs b/src/configs/hg_branch.rs
new file mode 100644
index 000000000..14e22636a
--- /dev/null
+++ b/src/configs/hg_branch.rs
@@ -0,0 +1,27 @@
+use crate::config::{ModuleConfig, RootModuleConfig, SegmentConfig};
+
+use ansi_term::{Color, Style};
+use starship_module_config_derive::ModuleConfig;
+
+#[derive(Clone, ModuleConfig)]
+pub struct HgBranchConfig<'a> {
+    pub symbol: SegmentConfig<'a>,
+    pub truncation_length: i64,
+    pub truncation_symbol: &'a str,
+    pub branch_name: SegmentConfig<'a>,
+    pub style: Style,
+    pub disabled: bool,
+}
+
+impl<'a> RootModuleConfig<'a> for HgBranchConfig<'a> {
+    fn new() -> Self {
+        HgBranchConfig {
+            symbol: SegmentConfig::new(" "),
+            truncation_length: std::i64::MAX,
+            truncation_symbol: "…",
+            branch_name: SegmentConfig::default(),
+            style: Color::Purple.bold(),
+            disabled: true,
+        }
+    }
+}
diff --git a/src/configs/mod.rs b/src/configs/mod.rs
index 51e4fa4dc..c30ac9a9d 100644
--- a/src/configs/mod.rs
+++ b/src/configs/mod.rs
@@ -10,6 +10,7 @@ pub mod git_branch;
 pub mod git_state;
 pub mod git_status;
 pub mod go;
+pub mod hg_branch;
 pub mod hostname;
 pub mod java;
 pub mod jobs;
diff --git a/src/configs/starship_root.rs b/src/configs/starship_root.rs
index 3873e96d1..eb0523a29 100644
--- a/src/configs/starship_root.rs
+++ b/src/configs/starship_root.rs
@@ -24,6 +24,7 @@ impl<'a> RootModuleConfig<'a> for StarshipRootConfig<'a> {
                 "git_branch",
                 "git_state",
                 "git_status",
+                "hg_branch",
                 "package",
                 // ↓ Toolchain version modules ↓
                 // (Let's keep these sorted alphabetically)
diff --git a/src/module.rs b/src/module.rs
index 82fdd03b6..7e3cb813b 100644
--- a/src/module.rs
+++ b/src/module.rs
@@ -21,6 +21,7 @@ pub const ALL_MODULES: &[&str] = &[
     "git_state",
     "git_status",
     "golang",
+    "hg_branch",
     "hostname",
     "java",
     "jobs",
diff --git a/src/modules/hg_branch.rs b/src/modules/hg_branch.rs
new file mode 100644
index 000000000..34b666c49
--- /dev/null
+++ b/src/modules/hg_branch.rs
@@ -0,0 +1,89 @@
+use std::process::Command;
+use unicode_segmentation::UnicodeSegmentation;
+
+use super::{Context, Module, RootModuleConfig};
+
+use crate::configs::hg_branch::HgBranchConfig;
+
+/// Creates a module with the Hg bookmark or branch in the current directory
+///
+/// Will display the bookmark or branch name if the current directory is an hg repo
+pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
+    let is_hg_repo = context
+        .try_begin_scan()?
+        .set_files(&[".hgignore"])
+        .set_folders(&[".hg"])
+        .is_match();
+
+    if !is_hg_repo {
+        return None;
+    }
+
+    let mut module = context.new_module("hg_branch");
+    let config = HgBranchConfig::try_load(module.config);
+    module.set_style(config.style);
+
+    module.get_prefix().set_value("on ");
+
+    let truncation_symbol = get_graphemes(config.truncation_symbol, 1);
+    module.create_segment("symbol", &config.symbol);
+
+    // TODO: Once error handling is implemented, warn the user if their config
+    // truncation length is nonsensical
+    let len = if config.truncation_length <= 0 {
+        log::warn!(
+            "\"truncation_length\" should be a positive value, found {}",
+            config.truncation_length
+        );
+        std::usize::MAX
+    } else {
+        config.truncation_length as usize
+    };
+
+    let get_branch_name = |tmpl| get_hg_log_template(tmpl, context);
+
+    let branch_name = get_branch_name("{activebookmark}")
+        .or_else(|| get_branch_name("{branch}"))
+        .unwrap_or_else(|| "(no branch)".to_string());
+
+    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.create_segment(
+        "name",
+        &config.branch_name.with_value(&truncated_and_symbol),
+    );
+
+    Some(module)
+}
+
+fn get_hg_log_template(hgtmpl: &str, ctx: &Context) -> Option<String> {
+    let output = Command::new("hg")
+        .args(&["log", "-r", ".", "--template", hgtmpl])
+        .current_dir(&ctx.current_dir)
+        .output()
+        .ok()
+        .and_then(|output| String::from_utf8(output.stdout).ok())?;
+
+    if output.is_empty() {
+        None
+    } else {
+        Some(output)
+    }
+}
+
+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/src/modules/mod.rs b/src/modules/mod.rs
index e32af5c88..2f891e8d3 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -10,6 +10,7 @@ mod git_branch;
 mod git_state;
 mod git_status;
 mod golang;
+mod hg_branch;
 mod hostname;
 mod java;
 mod jobs;
@@ -50,6 +51,7 @@ pub fn handle<'a>(module: &str, context: &'a Context) -> Option<Module<'a>> {
         "git_state" => git_state::module(context),
         "git_status" => git_status::module(context),
         "golang" => golang::module(context),
+        "hg_branch" => hg_branch::module(context),
         "hostname" => hostname::module(context),
         "java" => java::module(context),
         "jobs" => jobs::module(context),
diff --git a/tests/Dockerfile b/tests/Dockerfile
index af9c0332a..6b0af7d2a 100644
--- a/tests/Dockerfile
+++ b/tests/Dockerfile
@@ -61,6 +61,11 @@ RUN mkdir -p "$DOTNET_HOME" \
 ENV PATH $DOTNET_HOME:$PATH
 RUN dotnet help
 
+# Install Mercurial
+RUN HGPYTHON3=1 pip install mercurial
+# Check that Mercurial was correctly installed
+RUN hg --version
+
 # Create blank project
 RUN USER=nonroot cargo new --bin /src/starship
 WORKDIR /src/starship
diff --git a/tests/fixtures/hg-repo.bundle b/tests/fixtures/hg-repo.bundle
new file mode 100644
index 000000000..e41748507
--- /dev/null
+++ b/tests/fixtures/hg-repo.bundle
diff --git a/tests/testsuite/hg_branch.rs b/tests/testsuite/hg_branch.rs
new file mode 100644
index 000000000..12d2dc9bd
--- /dev/null
+++ b/tests/testsuite/hg_branch.rs
@@ -0,0 +1,219 @@
+use ansi_term::{Color, Style};
+use std::fs;
+use std::path::{Path, PathBuf};
+use std::process::Command;
+use std::{env, io};
+use tempfile;
+
+use crate::common::{self, TestCommand};
+
+enum Expect<'a> {
+    BranchName(&'a str),
+    Empty,
+    NoTruncation,
+    Symbol(&'a str),
+    Style(Style),
+    TruncationSymbol(&'a str),
+}
+
+#[test]
+#[ignore]
+fn test_hg_get_branch_fails() -> io::Result<()> {
+    let tempdir = tempfile::tempdir()?;
+
+    // Create a fake corrupted mercurial repo.
+    let hgdir = tempdir.path().join(".hg");
+    fs::create_dir(&hgdir)?;
+    fs::write(&hgdir.join("requires"), "fake-corrupted-repo")?;
+
+    expect_hg_branch_with_config(
+        tempdir.path(),
+        "",
+        &[Expect::BranchName(&"(no branch)"), Expect::NoTruncation],
+    )
+}
+
+#[test]
+#[ignore]
+fn test_hg_get_branch_autodisabled() -> io::Result<()> {
+    let tempdir = tempfile::tempdir()?;
+    expect_hg_branch_with_config(tempdir.path(), "", &[Expect::Empty])
+}
+
+#[test]
+#[ignore]
+fn test_hg_bookmark() -> io::Result<()> {
+    let tempdir = tempfile::tempdir()?;
+    let repo_dir = create_fixture_hgrepo(&tempdir)?;
+    run_hg(&["bookmark", "bookmark-101"], &repo_dir)?;
+    expect_hg_branch_with_config(
+        &repo_dir,
+        "",
+        &[Expect::BranchName(&"bookmark-101"), Expect::NoTruncation],
+    )
+}
+
+#[test]
+#[ignore]
+fn test_default_truncation_symbol() -> io::Result<()> {
+    let tempdir = tempfile::tempdir()?;
+    let repo_dir = create_fixture_hgrepo(&tempdir)?;
+    run_hg(&["branch", "-f", "branch-name-101"], &repo_dir)?;
+    run_hg(
+        &[
+            "commit",
+            "-m",
+            "empty commit 101",
+            "-u",
+            "fake user <fake@user>",
+        ],
+        &repo_dir,
+    )?;
+    expect_hg_branch_with_config(
+        &repo_dir,
+        "truncation_length = 14",
+        &[Expect::BranchName(&"branch-name-10")],
+    )
+}
+
+#[test]
+#[ignore]
+fn test_configured_symbols() -> io::Result<()> {
+    let tempdir = tempfile::tempdir()?;
+    let repo_dir = create_fixture_hgrepo(&tempdir)?;
+    run_hg(&["branch", "-f", "branch-name-121"], &repo_dir)?;
+    run_hg(
+        &[
+            "commit",
+            "-m",
+            "empty commit 121",
+            "-u",
+            "fake user <fake@user>",
+        ],
+        &repo_dir,
+    )?;
+    expect_hg_branch_with_config(
+        &repo_dir,
+        r#"
+          symbol = "B "
+          truncation_length = 14
+          truncation_symbol = "%"
+        "#,
+        &[
+            Expect::BranchName(&"branch-name-12"),
+            Expect::Symbol(&"B"),
+            Expect::TruncationSymbol(&"%"),
+        ],
+    )
+}
+
+#[test]
+#[ignore]
+fn test_configured_style() -> io::Result<()> {
+    let tempdir = tempfile::tempdir()?;
+    let repo_dir = create_fixture_hgrepo(&tempdir)?;
+    run_hg(&["branch", "-f", "branch-name-131"], &repo_dir)?;
+    run_hg(
+        &[
+            "commit",
+            "-m",
+            "empty commit 131",
+            "-u",
+            "fake user <fake@user>",
+        ],
+        &repo_dir,
+    )?;
+
+    expect_hg_branch_with_config(
+        &repo_dir,
+        r#"
+          style = "underline blue"
+        "#,
+        &[
+            Expect::BranchName(&"branch-name-131"),
+            Expect::Style(Color::Blue.underline()),
+            Expect::TruncationSymbol(&""),
+        ],
+    )
+}
+
+fn expect_hg_branch_with_config(
+    repo_dir: &Path,
+    config_options: &str,
+    expectations: &[Expect],
+) -> io::Result<()> {
+    let output = common::render_module("hg_branch")
+        .use_config(toml::from_str(&format!(
+            r#"
+            [hg_branch]
+            {}
+            "#,
+            config_options
+        ))?)
+        .arg("--path")
+        .arg(repo_dir.to_str().unwrap())
+        .output()?;
+
+    let actual = String::from_utf8(output.stdout).unwrap();
+
+    let mut expect_branch_name = "(no branch)";
+    let mut expect_style = Color::Purple.bold();
+    let mut expect_symbol = "\u{e0a0}";
+    let mut expect_truncation_symbol = "…";
+
+    for expect in expectations {
+        match expect {
+            Expect::Empty => {
+                assert_eq!("", actual);
+                return Ok(());
+            }
+            Expect::Symbol(symbol) => {
+                expect_symbol = symbol;
+            }
+            Expect::TruncationSymbol(truncation_symbol) => {
+                expect_truncation_symbol = truncation_symbol;
+            }
+            Expect::NoTruncation => {
+                expect_truncation_symbol = "";
+            }
+            Expect::BranchName(branch_name) => {
+                expect_branch_name = branch_name;
+            }
+            Expect::Style(style) => expect_style = *style,
+        }
+    }
+
+    let expected = format!(
+        "on {} ",
+        expect_style.paint(format!(
+            "{} {}{}",
+            expect_symbol, expect_branch_name, expect_truncation_symbol
+        )),
+    );
+    assert_eq!(expected, actual);
+    Ok(())
+}
+
+pub fn create_fixture_hgrepo(tempdir: &tempfile::TempDir) -> io::Result<PathBuf> {
+    let repo_path = tempdir.path().join("hg-repo");
+    let fixture_path = env::current_dir()?.join("tests/fixtures/hg-repo.bundle");
+
+    run_hg(
+        &[
+            "clone",
+            fixture_path.to_str().unwrap(),
+            repo_path.to_str().unwrap(),
+        ],
+        &tempdir.path(),
+    )?;
+
+    Ok(repo_path)
+}
+
+fn run_hg(args: &[&str], repo_dir: &Path) -> io::Result<()> {
+    Command::new("hg")
+        .args(args)
+        .current_dir(&repo_dir)
+        .output()?;
+    Ok(())
+}
diff --git a/tests/testsuite/main.rs b/tests/testsuite/main.rs
index 6efa3cfb8..c3e209262 100644
--- a/tests/testsuite/main.rs
+++ b/tests/testsuite/main.rs
@@ -11,6 +11,7 @@ mod git_branch;
 mod git_state;
 mod git_status;
 mod golang;
+mod hg_branch;
 mod hostname;
 mod jobs;
 mod line_break;