feat: Add the hg_branch module (#569)

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,
+        }
+    }
+}

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;

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)

src/module.rs

@@ -21,6 +21,7 @@ pub const ALL_MODULES: &[&str] = &[
     "git_state",
     "git_status",
     "golang",
+    "hg_branch",
     "hostname",
     "java",
     "jobs",

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()
+}

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),