From 5cdd8fd599734cdb4de9d664db975114b42ea69d Mon Sep 17 00:00:00 2001
From: Jason Shirk <jasonsh@microsoft.com>
Date: Tue, 5 May 2020 23:03:06 -0700
Subject: [PATCH] Add support for PowerShell (#65)

---
 CHANGELOG.md                            |  4 ++
 README.md                               | 12 ++++
 src/subcommand/init/mod.rs              |  2 +
 src/subcommand/init/shell/fish.rs       |  8 +--
 src/subcommand/init/shell/mod.rs        |  1 +
 src/subcommand/init/shell/powershell.rs | 85 +++++++++++++++++++++++++
 6 files changed, 108 insertions(+), 4 deletions(-)
 create mode 100644 src/subcommand/init/shell/powershell.rs

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d3d00a3..684894a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Added
+
+- Support for PowerShell.
+
 ### Removed
 
 - Backward compatibility with `v0.2.x` databases.
diff --git a/README.md b/README.md
index 7f8a86e..95cf7ae 100644
--- a/README.md
+++ b/README.md
@@ -15,6 +15,7 @@ A faster way to navigate your filesystem
     - [bash](#bash)
     - [fish](#fish)
     - [POSIX](#posix-shells)
+    - [PowerShell](#powershell)
     - [zsh](#zsh)
 - [Configuration](#configuration)
   - [`init` flags](#init-flags)
@@ -118,6 +119,17 @@ Add the following line to your `~/.zshrc`:
 eval "$(zoxide init zsh)"
 ```
 
+#### PowerShell
+
+Add the following line to your profile:
+
+```powershell
+Invoke-Expression (& {
+    $hook = if ($PSVersionTable.PSVersion.Major -lt 6) { 'prompt' } else { 'pwd' }
+    (zoxide init --hook $hook powershell) -join "`n"
+})
+```
+
 ## Configuration
 
 ### `init` flags
diff --git a/src/subcommand/init/mod.rs b/src/subcommand/init/mod.rs
index e8f6845..cf04cbd 100644
--- a/src/subcommand/init/mod.rs
+++ b/src/subcommand/init/mod.rs
@@ -43,6 +43,7 @@ impl Init {
             Shell::bash => shell::bash::CONFIG,
             Shell::fish => shell::fish::CONFIG,
             Shell::posix => shell::posix::CONFIG,
+            Shell::powershell => shell::powershell::CONFIG,
             Shell::zsh => shell::zsh::CONFIG,
         };
 
@@ -77,6 +78,7 @@ arg_enum! {
         bash,
         fish,
         posix,
+        powershell,
         zsh,
     }
 }
diff --git a/src/subcommand/init/shell/fish.rs b/src/subcommand/init/shell/fish.rs
index bd4b36f..c840b07 100644
--- a/src/subcommand/init/shell/fish.rs
+++ b/src/subcommand/init/shell/fish.rs
@@ -13,7 +13,7 @@ pub const CONFIG: ShellConfig = ShellConfig {
     },
 };
 
-fn z(z_cmd: &str) -> String {
+fn z(cmd: &str) -> String {
     format!(
         r#"
 function _z_cd
@@ -52,11 +52,11 @@ function {}
     end
 end
 "#,
-        z_cmd
+        cmd
     )
 }
 
-fn alias(z_cmd: &str) -> String {
+fn alias(cmd: &str) -> String {
     format!(
         r#"
 abbr -a {0}i '{0} -i'
@@ -69,7 +69,7 @@ abbr -a {0}qi 'zoxide query -i'
 abbr -a {0}r 'zoxide remove'
 abbr -a {0}ri 'zoxide remove -i'
 "#,
-        z_cmd
+        cmd
     )
 }
 
diff --git a/src/subcommand/init/shell/mod.rs b/src/subcommand/init/shell/mod.rs
index 01172e7..45413e3 100644
--- a/src/subcommand/init/shell/mod.rs
+++ b/src/subcommand/init/shell/mod.rs
@@ -1,6 +1,7 @@
 pub mod bash;
 pub mod fish;
 pub mod posix;
+pub mod powershell;
 pub mod zsh;
 
 use anyhow::Result;
diff --git a/src/subcommand/init/shell/powershell.rs b/src/subcommand/init/shell/powershell.rs
new file mode 100644
index 0000000..142dd7e
--- /dev/null
+++ b/src/subcommand/init/shell/powershell.rs
@@ -0,0 +1,85 @@
+use super::{HookConfig, ShellConfig};
+
+use anyhow::Result;
+
+use std::borrow::Cow;
+
+pub const CONFIG: ShellConfig = ShellConfig {
+    z,
+    alias,
+    hook: HookConfig {
+        prompt: HOOK_PROMPT,
+        pwd: hook_pwd,
+    },
+};
+
+fn z(cmd: &str) -> String {
+    format!(
+        r#"
+function {} {{
+    function z_cd($dir) {{
+        Set-Location $dir -ea Stop
+        if ($env:_ZO_ECHO -eq "1") {{
+            Write-Host "$PWD"
+        }}
+    }}
+
+    if ($args.Length -eq 0) {{
+        z_cd ~
+    }}
+    elseif ($args.Length -eq 1 -and $args[0] -eq '-') {{
+        z_cd -
+    }}
+    else {{
+        $result = zoxide query @args
+        if ($LASTEXITCODE -eq 0 -and $result -is [string] -and (Test-Path $result)) {{
+            z_cd $result
+        }} else {{
+            $result
+        }}
+    }}
+}}
+"#,
+        cmd
+    )
+}
+
+fn alias(cmd: &str) -> String {
+    format!(
+        r#"
+function {0}i {{ {0} -i @args }}
+
+function {0}a {{ zoxide add @args }}
+
+function {0}q {{ zoxide query @args }}
+function {0}qi {{ zoxide query -i @args }}
+
+function {0}r {{ zoxide remove @args }}
+function {0}ri {{ zoxide remove -i @args }}
+"#,
+        cmd
+    )
+}
+
+const HOOK_PROMPT: &str = r#"
+$PreZoxidePrompt = $function:prompt
+function prompt {
+    $null = zoxide add
+    & $PreZoxidePrompt
+}
+"#;
+
+const fn hook_pwd() -> Result<Cow<'static, str>> {
+    const HOOK_PWD: &str = r#"
+if ($PSVersionTable.PSVersion.Major -ge 6) {
+    $ExecutionContext.InvokeCommand.LocationChangedAction = {
+        $null = zoxide add
+    }
+} else {
+    Write-Error "pwd hook requires pwsh - use 'zoxide init powershell --hook prompt'"
+}
+"#;
+
+    Ok(Cow::Borrowed(HOOK_PWD))
+}
+