feat: add mercurial state (#6745)

.github/config-schema.json

@@ -892,6 +892,26 @@
         }
       ]
     },
+    "hg_state": {
+      "default": {
+        "bisect": "BISECTING",
+        "disabled": true,
+        "format": "\\([$state]($style)\\) ",
+        "graft": "GRAFTING",
+        "histedit": "HISTEDITING",
+        "merge": "MERGING",
+        "rebase": "REBASING",
+        "shelve": "SHELVING",
+        "style": "bold yellow",
+        "transplant": "TRANSPLANTING",
+        "update": "UPDATING"
+      },
+      "allOf": [
+        {
+          "$ref": "#/definitions/HgStateConfig"
+        }
+      ]
+    },
     "hostname": {
       "default": {
         "aliases": {},
@@ -4232,6 +4252,56 @@
       },
       "additionalProperties": false
     },
+    "HgStateConfig": {
+      "type": "object",
+      "properties": {
+        "merge": {
+          "default": "MERGING",
+          "type": "string"
+        },
+        "rebase": {
+          "default": "REBASING",
+          "type": "string"
+        },
+        "update": {
+          "default": "UPDATING",
+          "type": "string"
+        },
+        "bisect": {
+          "default": "BISECTING",
+          "type": "string"
+        },
+        "shelve": {
+          "default": "SHELVING",
+          "type": "string"
+        },
+        "graft": {
+          "default": "GRAFTING",
+          "type": "string"
+        },
+        "transplant": {
+          "default": "TRANSPLANTING",
+          "type": "string"
+        },
+        "histedit": {
+          "default": "HISTEDITING",
+          "type": "string"
+        },
+        "style": {
+          "default": "bold yellow",
+          "type": "string"
+        },
+        "format": {
+          "default": "\\([$state]($style)\\) ",
+          "type": "string"
+        },
+        "disabled": {
+          "default": true,
+          "type": "boolean"
+        }
+      },
+      "additionalProperties": false
+    },
     "HostnameConfig": {
       "type": "object",
       "properties": {

docs/config/README.md

@@ -281,6 +281,7 @@ $git_state\
 $git_metrics\
 $git_status\
 $hg_branch\
+$hg_state\
 $pijul_channel\
 $docker_context\
 $package\
@@ -2990,6 +2991,39 @@ truncation_length = 4
 truncation_symbol = ''
 ```
 
+## Mercurial State
+
+The `hg_state` module will show in directories which are part of a mercurial
+repository, and where there is an operation in progress, such as: _REBASING_,
+_BISECTING_, etc.
+
+### Options
+
+| Option       | Default                   | Description                                                   |
+| ------------ | ------------------------- | ------------------------------------------------------------- |
+| `merge`      | `'MERGING'`               | A format string displayed when a `merge` is in progress.      |
+| `rebase`     | `'REBASING'`              | A format string displayed when a `rebase` is in progress.     |
+| `update`     | `'UPDATING'`              | A format string displayed when a `update` is in progress.     |
+| `bisect`     | `'BISECTING'`             | A format string displayed when a `bisect` is in progress.     |
+| `shelve`     | `'SHELVING'`              | A format string displayed when a `shelve` is in progress.     |
+| `graft`      | `'GRAFTING'`              | A format string displayed when a `graft` is in progress.      |
+| `transplant` | `'TRANSPLANTING'`         | A format string displayed when a `transplant` is in progress. |
+| `histedit`   | `'HISTEDITING'`           | A format string displayed when a `histedit` is in progress.   |
+| `style`      | `'bold yellow'`           | The style for the module.                                     |
+| `format`     | `'\([$state]($style)\) '` | The format for the module.                                    |
+| `disabled`   | `true`                    | Disables the `hg_state` module.                               |
+
+### Variables
+
+| Variable         | Example    | Description                         |
+| ---------------- | ---------- | ----------------------------------- |
+| state            | `REBASING` | The current state of the repo       |
+| progress_current | `1`        | The current operation progress      |
+| progress_total   | `2`        | The total operation progress        |
+| style\*          |            | Mirrors the value of option `style` |
+
+*: This variable can only be used as a part of a style string
+
 ## Mise
 
 The `mise` module shows the current mise health as reported by running `mise doctor`.

docs/public/presets/toml/bracketed-segments.toml

@@ -85,6 +85,9 @@ format = '\[[$symbol($version)]($style)\]'
 [hg_branch]
 format = '\[[$symbol$branch]($style)\]'
 
+[hg_state]
+format = '\([$state]($style)\) '
+
 [java]
 format = '\[[$symbol($version)]($style)\]'
 

src/configs/hg_state.rs

@@ -0,0 +1,40 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone, Deserialize, Serialize)]
+#[cfg_attr(
+    feature = "config-schema",
+    derive(schemars::JsonSchema),
+    schemars(deny_unknown_fields)
+)]
+#[serde(default)]
+pub struct HgStateConfig<'a> {
+    pub merge: &'a str,
+    pub rebase: &'a str,
+    pub update: &'a str,
+    pub bisect: &'a str,
+    pub shelve: &'a str,
+    pub graft: &'a str,
+    pub transplant: &'a str,
+    pub histedit: &'a str,
+    pub style: &'a str,
+    pub format: &'a str,
+    pub disabled: bool,
+}
+
+impl Default for HgStateConfig<'_> {
+    fn default() -> Self {
+        HgStateConfig {
+            merge: "MERGING",
+            rebase: "REBASING",
+            update: "UPDATING",
+            bisect: "BISECTING",
+            shelve: "SHELVING",
+            graft: "GRAFTING",
+            transplant: "TRANSPLANTING",
+            histedit: "HISTEDITING",
+            style: "bold yellow",
+            format: "\\([$state]($style)\\) ",
+            disabled: true,
+        }
+    }
+}

src/configs/mod.rs

@@ -46,6 +46,7 @@ pub mod haskell;
 pub mod haxe;
 pub mod helm;
 pub mod hg_branch;
+pub mod hg_state;
 pub mod hostname;
 pub mod java;
 pub mod jobs;
@@ -206,6 +207,8 @@ pub struct FullConfig<'a> {
     #[serde(borrow)]
     hg_branch: hg_branch::HgBranchConfig<'a>,
     #[serde(borrow)]
+    hg_state: hg_state::HgStateConfig<'a>,
+    #[serde(borrow)]
     hostname: hostname::HostnameConfig<'a>,
     #[serde(borrow)]
     java: java::JavaConfig<'a>,

src/configs/starship_root.rs

@@ -48,6 +48,7 @@ pub const PROMPT_ORDER: &[&str] = &[
     "git_metrics",
     "git_status",
     "hg_branch",
+    "hg_state",
     "pijul_channel",
     "docker_context",
     "package",

src/module.rs

@@ -50,6 +50,7 @@ pub const ALL_MODULES: &[&str] = &[
     "haxe",
     "helm",
     "hg_branch",
+    "hg_state",
     "hostname",
     "java",
     "jobs",

src/modules/hg_state.rs

@@ -0,0 +1,277 @@
+use std::fs::File;
+use std::io::{self, Read};
+use std::path::Path;
+
+use super::{Context, Module, ModuleConfig};
+use crate::configs::hg_state::HgStateConfig;
+use crate::formatter::StringFormatter;
+
+/// Creates a module with the state of hg repository at the current directory
+///
+/// During a mercurial operation it will show: MERGING, REBASING, UPDATING etc.
+pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
+    let mut module = context.new_module("hg_state");
+    let config: HgStateConfig = HgStateConfig::try_load(module.config);
+
+    // As we default to disabled=true, we have to check here after loading our config module,
+    // before it was only checking against whatever is in the config starship.toml
+    if config.disabled {
+        return None;
+    };
+
+    let repo_root = context.begin_ancestor_scan().set_folders(&[".hg"]).scan()?;
+
+    let state_description = get_state_description(&repo_root, &config)?;
+
+    let parsed = StringFormatter::new(config.format).and_then(|formatter| {
+        formatter
+            .map_meta(|variable, _| match variable {
+                "state" => Some(state_description.label),
+                _ => None,
+            })
+            .map_style(|variable| match variable {
+                "style" => Some(Ok(config.style)),
+                _ => None,
+            })
+            .parse(None, Some(context))
+    });
+
+    module.set_segments(match parsed {
+        Ok(segments) => segments,
+        Err(error) => {
+            log::warn!("Error in module `hg_state`:\n{}", error);
+            return None;
+        }
+    });
+
+    Some(module)
+}
+
+fn get_state_description<'a>(
+    hg_root: &Path,
+    config: &HgStateConfig<'a>,
+) -> Option<StateDescription<'a>> {
+    if hg_root.join(".hg").join("rebasestate").exists() {
+        Some(StateDescription {
+            label: config.rebase,
+        })
+    } else if hg_root.join(".hg").join("updatestate").exists() {
+        Some(StateDescription {
+            label: config.update,
+        })
+    } else if hg_root.join(".hg").join("bisect.state").exists() {
+        Some(StateDescription {
+            label: config.bisect,
+        })
+    } else if hg_root.join(".hg").join("graftstate").exists() {
+        Some(StateDescription {
+            label: config.graft,
+        })
+    } else if hg_root
+        .join(".hg")
+        .join("transplant")
+        .join("journal")
+        .exists()
+    {
+        Some(StateDescription {
+            label: config.transplant,
+        })
+    } else if hg_root.join(".hg").join("histedit-state").exists() {
+        Some(StateDescription {
+            label: config.histedit,
+        })
+    } else if is_merge_state(hg_root).is_ok() {
+        Some(StateDescription {
+            label: config.merge,
+        })
+    } else {
+        return None;
+    }
+}
+
+fn is_merge_state(hg_root: &Path) -> io::Result<bool> {
+    let dirstate_path = hg_root.join(".hg").join("dirstate");
+
+    let mut file = File::open(dirstate_path)?;
+    let mut buffer = [0u8; 40]; // First 40 bytes: 20 for p1, 20 for p2
+    file.read_exact(&mut buffer)?;
+
+    let p2 = &buffer[20..40];
+    let is_merge = p2.iter().any(|&b| b != 0);
+
+    Ok(is_merge)
+}
+
+struct StateDescription<'a> {
+    label: &'a str,
+}
+
+#[cfg(test)]
+mod tests {
+    use nu_ansi_term::Color;
+    use std::fs::{File, create_dir_all};
+    use std::io;
+    use std::path::Path;
+
+    use crate::test::{FixtureProvider, ModuleRenderer, fixture_repo};
+    use crate::utils::create_command;
+
+    #[test]
+    fn test_hg_show_nothing_on_empty_dir() -> io::Result<()> {
+        let repo_dir = tempfile::tempdir()?;
+
+        let actual = ModuleRenderer::new("hg_state")
+            .path(repo_dir.path())
+            .collect();
+
+        let expected = None;
+
+        assert_eq!(expected, actual);
+        repo_dir.close()
+    }
+
+    #[test]
+    #[ignore]
+    fn test_hg_disabled_per_default() -> io::Result<()> {
+        let tempdir = fixture_repo(FixtureProvider::Hg)?;
+        let repo_dir = tempdir.path();
+        run_hg(&["whatever", "blubber"], repo_dir)?;
+        let actual = ModuleRenderer::new("hg_state")
+            .path(repo_dir.to_str().unwrap())
+            .collect();
+
+        let expected = None;
+
+        assert_eq!(expected, actual);
+        tempdir.close()
+    }
+
+    #[test]
+    #[ignore]
+    fn test_hg_merging() -> io::Result<()> {
+        let tempdir = fixture_repo(FixtureProvider::Hg)?;
+        let repo_dir = tempdir.path();
+        run_hg(&["branch", "-f", "branch-name-101"], repo_dir)?;
+        run_hg(
+            &[
+                "commit",
+                "-m",
+                "empty commit 101",
+                "-u",
+                "fake user <fake@user>",
+            ],
+            repo_dir,
+        )?;
+        run_hg(&["update", "-r", "branch-name-0"], repo_dir)?;
+        run_hg(&["merge", "-r", "branch-name-101"], repo_dir)?;
+        expect_hg_state_with_config(
+            repo_dir,
+            Some(format!("({}) ", Color::Yellow.bold().paint("MERGING"))),
+        );
+        tempdir.close()
+    }
+
+    #[test]
+    #[ignore]
+    fn test_hg_rebasing() -> io::Result<()> {
+        let tempdir = fixture_repo(FixtureProvider::Hg)?;
+        let repo_dir = tempdir.path();
+        File::create(repo_dir.join(".hg").join("rebasestate"))?;
+        expect_hg_state_with_config(
+            repo_dir,
+            Some(format!("({}) ", Color::Yellow.bold().paint("REBASING"))),
+        );
+        tempdir.close()
+    }
+
+    #[test]
+    #[ignore]
+    fn test_hg_updating() -> io::Result<()> {
+        let tempdir = fixture_repo(FixtureProvider::Hg)?;
+        let repo_dir = tempdir.path();
+        File::create(repo_dir.join(".hg").join("updatestate"))?;
+        expect_hg_state_with_config(
+            repo_dir,
+            Some(format!("({}) ", Color::Yellow.bold().paint("UPDATING"))),
+        );
+        tempdir.close()
+    }
+
+    #[test]
+    #[ignore]
+    fn test_hg_bisecting() -> io::Result<()> {
+        let tempdir = fixture_repo(FixtureProvider::Hg)?;
+        let repo_dir = tempdir.path();
+        run_hg(&["bisect", "--good"], repo_dir)?;
+        expect_hg_state_with_config(
+            repo_dir,
+            Some(format!("({}) ", Color::Yellow.bold().paint("BISECTING"))),
+        );
+        tempdir.close()
+    }
+
+    #[test]
+    #[ignore]
+    fn test_hg_grafting() -> io::Result<()> {
+        let tempdir = fixture_repo(FixtureProvider::Hg)?;
+        let repo_dir = tempdir.path();
+        File::create(repo_dir.join(".hg").join("graftstate"))?;
+        expect_hg_state_with_config(
+            repo_dir,
+            Some(format!("({}) ", Color::Yellow.bold().paint("GRAFTING"))),
+        );
+        tempdir.close()
+    }
+
+    #[test]
+    #[ignore]
+    fn test_hg_transplanting() -> io::Result<()> {
+        let tempdir = fixture_repo(FixtureProvider::Hg)?;
+        let repo_dir = tempdir.path();
+        create_dir_all(repo_dir.join(".hg").join("transplant"))?;
+        File::create(repo_dir.join(".hg").join("transplant").join("journal"))?;
+        expect_hg_state_with_config(
+            repo_dir,
+            Some(format!(
+                "({}) ",
+                Color::Yellow.bold().paint("TRANSPLANTING")
+            )),
+        );
+        tempdir.close()
+    }
+
+    #[test]
+    #[ignore]
+    fn test_hg_histediting() -> io::Result<()> {
+        let tempdir = fixture_repo(FixtureProvider::Hg)?;
+        let repo_dir = tempdir.path();
+        File::create(repo_dir.join(".hg").join("histedit-state"))?;
+        expect_hg_state_with_config(
+            repo_dir,
+            Some(format!("({}) ", Color::Yellow.bold().paint("HISTEDITING"))),
+        );
+        tempdir.close()
+    }
+
+    fn run_hg(args: &[&str], repo_dir: &Path) -> io::Result<()> {
+        create_command("hg")?
+            .args(args)
+            .current_dir(repo_dir)
+            .output()?;
+        Ok(())
+    }
+
+    fn expect_hg_state_with_config(repo_dir: &Path, expected: Option<String>) {
+        let actual = ModuleRenderer::new("hg_state")
+            .path(repo_dir.to_str().unwrap())
+            .config({
+                toml::toml! {
+                    [hg_state]
+                    disabled = false
+                }
+            })
+            .collect();
+
+        assert_eq!(expected, actual);
+    }
+}

src/modules/mod.rs

@@ -43,6 +43,7 @@ mod haskell;
 mod haxe;
 mod helm;
 mod hg_branch;
+mod hg_state;
 mod hostname;
 mod java;
 mod jobs;
@@ -159,6 +160,7 @@ pub fn handle<'a>(module: &str, context: &'a Context) -> Option<Module<'a>> {
             "haxe" => haxe::module(context),
             "helm" => helm::module(context),
             "hg_branch" => hg_branch::module(context),
+            "hg_state" => hg_state::module(context),
             "hostname" => hostname::module(context),
             "java" => java::module(context),
             "jobs" => jobs::module(context),
@@ -288,6 +290,7 @@ pub fn description(module: &str) -> &'static str {
         "haxe" => "The currently installed version of Haxe",
         "helm" => "The currently installed version of Helm",
         "hg_branch" => "The active branch and topic of the repo in your current directory",
+        "hg_state" => "The current hg operation",
         "hostname" => "The system hostname",
         "java" => "The currently installed version of Java",
         "jobs" => "The current number of jobs running",