From 318050976b58f032a17ff4d122d21cb63faf8e3b Mon Sep 17 00:00:00 2001
From: Nemo157 <git@nemo157.com>
Date: Mon, 12 Jun 2023 11:05:44 +0200
Subject: [PATCH] feat(pwsh): Support vi command mode indicator (#5049)

Support vi command mode in powershell
---
 src/init/starship.ps1    | 10 ++++++++++
 src/modules/character.rs | 38 +++++++++++++++++++++++++++++++++++---
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/src/init/starship.ps1 b/src/init/starship.ps1
index 4422a1ed2bce..f31d223db042 100755
--- a/src/init/starship.ps1
+++ b/src/init/starship.ps1
@@ -141,6 +141,10 @@ $null = New-Module starship {
 
         $arguments += "--status=$($lastExitCodeForPrompt)"
 
+        if ([Microsoft.PowerShell.PSConsoleReadLine]::InViCommandMode()) {
+            $arguments += "--keymap=vi"
+        }
+
         # Invoke Starship
         $promptText = if ($script:TransientPrompt) {
             $script:TransientPrompt = $false
@@ -206,6 +210,12 @@ $null = New-Module starship {
         )
     )
 
+    try {
+        Set-PSReadLineOption -ViModeIndicator script -ViModeChangeHandler {
+            [Microsoft.PowerShell.PSConsoleReadLine]::InvokePrompt()
+        }
+    } catch {}
+
     Export-ModuleMember -Function @(
         "Enable-TransientPrompt"
         "Disable-TransientPrompt"
diff --git a/src/modules/character.rs b/src/modules/character.rs
index 7fd19c8ecabc..c897b42601ca 100644
--- a/src/modules/character.rs
+++ b/src/modules/character.rs
@@ -35,9 +35,9 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     // We do some environment detection in src/init.rs to translate.
     // The result: in non-vi fish, keymap is always reported as "insert"
     let mode = match (&context.shell, keymap) {
-        (Shell::Fish, "default") | (Shell::Zsh, "vicmd") | (Shell::Cmd, "vi") => {
-            ShellEditMode::Normal
-        }
+        (Shell::Fish, "default")
+        | (Shell::Zsh, "vicmd")
+        | (Shell::Cmd | Shell::PowerShell, "vi") => ShellEditMode::Normal,
         (Shell::Fish, "visual") => ShellEditMode::Visual,
         (Shell::Fish, "replace") => ShellEditMode::Replace,
         (Shell::Fish, "replace_one") => ShellEditMode::ReplaceOne,
@@ -260,4 +260,36 @@ mod test {
             .collect();
         assert_eq!(expected_other, actual);
     }
+
+    #[test]
+    fn powershell_keymap() {
+        let expected_vicmd = Some(format!("{} ", Color::Green.bold().paint("❮")));
+        let expected_specified = Some(format!("{} ", Color::Green.bold().paint("V")));
+        let expected_other = Some(format!("{} ", Color::Green.bold().paint("❯")));
+
+        // powershell keymap is vi
+        let actual = ModuleRenderer::new("character")
+            .shell(Shell::PowerShell)
+            .keymap("vi")
+            .collect();
+        assert_eq!(expected_vicmd, actual);
+
+        // specified vicmd character
+        let actual = ModuleRenderer::new("character")
+            .config(toml::toml! {
+                [character]
+                vicmd_symbol = "[V](bold green)"
+            })
+            .shell(Shell::PowerShell)
+            .keymap("vi")
+            .collect();
+        assert_eq!(expected_specified, actual);
+
+        // powershell keymap is other
+        let actual = ModuleRenderer::new("character")
+            .shell(Shell::PowerShell)
+            .keymap("visual")
+            .collect();
+        assert_eq!(expected_other, actual);
+    }
 }