From 78f33bc002816dd690047a7a38077b6ac6617d44 Mon Sep 17 00:00:00 2001
From: "John W. Parent" <45471568+johnwparent@users.noreply.github.com>
Date: Tue, 27 Jun 2023 21:26:51 -0400
Subject: [PATCH] Windows: Add PowerShell env support (#37951)

PowerShell requires explicit shell and env support in Spack.
This is due to the distinct differences in shell interactions between
cmd and pwsh. Add a doskey in pwsh piping 'spack' commands to a
powershell script similar to the sh function 'spack'. Add
support for PowerShell-specific shell interactions from Spack
(set/unset shell variables).
---
 bin/spack.bat                            |   2 +-
 bin/spack.ps1                            | 132 +++++++++++++++++++++++
 lib/spack/spack/cmd/common/__init__.py   |   6 +-
 lib/spack/spack/cmd/env.py               |   7 ++
 lib/spack/spack/environment/shell.py     |   4 +
 lib/spack/spack/test/cmd/build_env.py    |   7 +-
 lib/spack/spack/test/conftest.py         |  18 ++++
 lib/spack/spack/test/util/environment.py |  22 ++--
 lib/spack/spack/util/environment.py      |   8 +-
 share/spack/setup-env.ps1                |   4 +
 share/spack/spack-completion.bash        |   2 +-
 11 files changed, 197 insertions(+), 15 deletions(-)
 create mode 100644 bin/spack.ps1

diff --git a/bin/spack.bat b/bin/spack.bat
index 514c2706d7..9aff863a90 100644
--- a/bin/spack.bat
+++ b/bin/spack.bat
@@ -214,7 +214,7 @@ goto :end_switch
 if defined _sp_args (
     if NOT "%_sp_args%"=="%_sp_args:--help=%" (
         goto :default_case
-    ) else if NOT "%_sp_args%"=="%_sp_args: -h=%" (
+    ) else if NOT "%_sp_args%"=="%_sp_args:-h=%" (
         goto :default_case
     ) else if NOT "%_sp_args%"=="%_sp_args:--bat=%" (
         goto :default_case
diff --git a/bin/spack.ps1 b/bin/spack.ps1
new file mode 100644
index 0000000000..39fe0167ca
--- /dev/null
+++ b/bin/spack.ps1
@@ -0,0 +1,132 @@
+#  Copyright 2013-2023 Lawrence Livermore National Security, LLC and other
+#  Spack Project Developers. See the top-level COPYRIGHT file for details.
+
+#  SPDX-License-Identifier: (Apache-2.0 OR MIT)
+# #######################################################################
+
+function Compare-CommonArgs {
+    $CMDArgs = $args[0]
+    # These aruments take precedence and call for no futher parsing of arguments
+    # invoke actual Spack entrypoint with that context and exit after
+    "--help", "-h", "--version", "-V" | ForEach-Object {
+        $arg_opt = $_
+        if(($CMDArgs) -and ([bool]($CMDArgs.Where({$_ -eq $arg_opt})))) {
+            return $true
+        }
+    }
+    return $false
+}
+
+function Read-SpackArgs {
+    $SpackCMD_params = @()
+    $SpackSubCommand = $NULL
+    $SpackSubCommandArgs = @()
+    $args_ = $args[0]
+    $args_ | ForEach-Object {
+        if (!$SpackSubCommand) {
+            if($_.SubString(0,1) -eq "-")
+            {
+                $SpackCMD_params += $_
+            }
+            else{
+                $SpackSubCommand = $_
+            }
+        }
+        else{
+            $SpackSubCommandArgs += $_
+        }
+    }
+    return $SpackCMD_params, $SpackSubCommand, $SpackSubCommandArgs
+}
+
+function Invoke-SpackCD {
+    if (Compare-CommonArgs $SpackSubCommandArgs) {
+        python $Env:SPACK_ROOT/bin/spack cd -h
+    }
+    else {
+        $LOC = $(python $Env:SPACK_ROOT/bin/spack location $SpackSubCommandArgs)
+        if (($NULL -ne $LOC)){
+            if ( Test-Path -Path $LOC){
+                Set-Location $LOC
+            }
+            else{
+                exit 1
+            }
+        }
+        else {
+            exit 1
+        }
+    }
+}
+
+function Invoke-SpackEnv {
+    if (Compare-CommonArgs $SpackSubCommandArgs[0]) {
+        python $Env:SPACK_ROOT/bin/spack env -h
+    }
+    else {
+        $SubCommandSubCommand = $SpackSubCommandArgs[0]
+        $SubCommandSubCommandArgs = $SpackSubCommandArgs[1..$SpackSubCommandArgs.Count]
+        switch ($SubCommandSubCommand) {
+            "activate" {
+                if (Compare-CommonArgs $SubCommandSubCommandArgs) {
+                    python $Env:SPACK_ROOT/bin/spack env activate $SubCommandSubCommandArgs
+                }
+                elseif ([bool]($SubCommandSubCommandArgs.Where({$_ -eq "--pwsh"}))) {
+                    python $Env:SPACK_ROOT/bin/spack env activate $SubCommandSubCommandArgs
+                }
+                elseif (!$SubCommandSubCommandArgs) {
+                    python $Env:SPACK_ROOT/bin/spack env activate $SubCommandSubCommandArgs
+                }
+                else {
+                    $SpackEnv = $(python $Env:SPACK_ROOT/bin/spack $SpackCMD_params env activate "--pwsh" $SubCommandSubCommandArgs)
+                    $ExecutionContext.InvokeCommand($SpackEnv)
+                }
+            }
+            "deactivate" {
+                if ([bool]($SubCommandSubCommandArgs.Where({$_ -eq "--pwsh"}))) {
+                    python $Env:SPACK_ROOT/bin/spack env deactivate $SubCommandSubCommandArgs
+                }
+                elseif($SubCommandSubCommandArgs) {
+                    python $Env:SPACK_ROOT/bin/spack env deactivate -h
+                }
+                else {
+                    $SpackEnv = $(python $Env:SPACK_ROOT/bin/spack $SpackCMD_params env deactivate --pwsh)
+                    $ExecutionContext.InvokeCommand($SpackEnv)
+                }
+            }
+            default {python $Env:SPACK_ROOT/bin/spack $SpackCMD_params $SpackSubCommand $SpackSubCommandArgs}
+        }
+    }
+}
+
+function Invoke-SpackLoad {
+    if (Compare-CommonArgs $SpackSubCommandArgs) {
+        python $Env:SPACK_ROOT/bin/spack $SpackCMD_params $SpackSubCommand $SpackSubCommandArgs
+    }
+    elseif ([bool]($SpackSubCommandArgs.Where({($_ -eq "--pwsh") -or ($_ -eq "--list")}))) {
+        python $Env:SPACK_ROOT/bin/spack $SpackCMD_params $SpackSubCommand $SpackSubCommandArgs
+    }
+    else {
+        $SpackEnv = $(python $Env:SPACK_ROOT/bin/spack $SpackCMD_params $SpackSubCommand "--pwsh" $SpackSubCommandArgs)
+        $ExecutionContext.InvokeCommand($SpackEnv)
+    }
+}
+
+
+$SpackCMD_params, $SpackSubCommand, $SpackSubCommandArgs = Read-SpackArgs $args
+
+if (Compare-CommonArgs $SpackCMD_params) {
+    python $Env:SPACK_ROOT/bin/spack $SpackCMD_params $SpackSubCommand $SpackSubCommandArgs
+    exit $LASTEXITCODE
+}
+
+# Process Spack commands with special conditions
+# all other commands are piped directly to Spack
+switch($SpackSubCommand)
+{
+    "cd"     {Invoke-SpackCD}
+    "env"    {Invoke-SpackEnv}
+    "load"   {Invoke-SpackLoad}
+    "unload" {Invoke-SpackLoad}
+    default  {python $Env:SPACK_ROOT/bin/spack $SpackCMD_params $SpackSubCommand $SpackSubCommandArgs}
+}
diff --git a/lib/spack/spack/cmd/common/__init__.py b/lib/spack/spack/cmd/common/__init__.py
index 28f4f291a2..51329a1b6c 100644
--- a/lib/spack/spack/cmd/common/__init__.py
+++ b/lib/spack/spack/cmd/common/__init__.py
@@ -36,7 +36,10 @@ def shell_init_instructions(cmd, equivalent):
         "  source %s/setup-env.fish" % spack.paths.share_path,
         "",
         color.colorize("@*c{For Windows batch:}"),
-        "  source %s/spack_cmd.bat" % spack.paths.share_path,
+        "  %s\\spack_cmd.bat" % spack.paths.bin_path,
+        "",
+        color.colorize("@*c{For PowerShell:}"),
+        "  %s\\setup-env.ps1" % spack.paths.share_path,
         "",
         "Or, if you do not want to use shell support, run "
         + ("one of these" if shell_specific else "this")
@@ -50,6 +53,7 @@ def shell_init_instructions(cmd, equivalent):
             equivalent.format(sh_arg="--csh ") + "  # csh/tcsh",
             equivalent.format(sh_arg="--fish") + "  # fish",
             equivalent.format(sh_arg="--bat ") + "  # batch",
+            equivalent.format(sh_arg="--pwsh") + "  # powershell",
         ]
     else:
         msg += ["  " + equivalent]
diff --git a/lib/spack/spack/cmd/env.py b/lib/spack/spack/cmd/env.py
index 6774b15afd..93e22031d8 100644
--- a/lib/spack/spack/cmd/env.py
+++ b/lib/spack/spack/cmd/env.py
@@ -86,6 +86,13 @@ def env_activate_setup_parser(subparser):
         const="bat",
         help="print bat commands to activate the environment",
     )
+    shells.add_argument(
+        "--pwsh",
+        action="store_const",
+        dest="shell",
+        const="pwsh",
+        help="print powershell commands to activate environment",
+    )
 
     view_options = subparser.add_mutually_exclusive_group()
     view_options.add_argument(
diff --git a/lib/spack/spack/environment/shell.py b/lib/spack/spack/environment/shell.py
index 4df66ac28c..c6d2b06c98 100644
--- a/lib/spack/spack/environment/shell.py
+++ b/lib/spack/spack/environment/shell.py
@@ -42,6 +42,8 @@ def activate_header(env, shell, prompt=None):
         cmds += 'set "SPACK_ENV=%s"\n' % env.path
         # TODO: despacktivate
         # TODO: prompt
+    elif shell == "pwsh":
+        cmds += "$Env:SPACK_ENV=%s\n" % env.path
     else:
         if "color" in os.getenv("TERM", "") and prompt:
             prompt = colorize("@G{%s}" % prompt, color=True, enclose=True)
@@ -79,6 +81,8 @@ def deactivate_header(shell):
         cmds += 'set "SPACK_ENV="\n'
         # TODO: despacktivate
         # TODO: prompt
+    elif shell == "pwsh":
+        cmds += "Remove-Item Env:SPACK_ENV"
     else:
         cmds += "if [ ! -z ${SPACK_ENV+x} ]; then\n"
         cmds += "unset SPACK_ENV; export SPACK_ENV;\n"
diff --git a/lib/spack/spack/test/cmd/build_env.py b/lib/spack/spack/test/cmd/build_env.py
index a03d9760bf..c27b292cb1 100644
--- a/lib/spack/spack/test/cmd/build_env.py
+++ b/lib/spack/spack/test/cmd/build_env.py
@@ -35,12 +35,15 @@ def test_build_env_requires_a_spec(args):
 _out_file = "env.out"
 
 
+@pytest.mark.parametrize("shell", ["pwsh", "bat"] if sys.platform == "win32" else ["bash"])
 @pytest.mark.usefixtures("config", "mock_packages", "working_env")
-def test_dump(tmpdir):
+def test_dump(shell_as, shell, tmpdir):
     with tmpdir.as_cwd():
         build_env("--dump", _out_file, "zlib")
         with open(_out_file) as f:
-            if sys.platform == "win32":
+            if shell == "pwsh":
+                assert any(line.startswith("$Env:PATH") for line in f.readlines())
+            elif shell == "bat":
                 assert any(line.startswith('set "PATH=') for line in f.readlines())
             else:
                 assert any(line.startswith("PATH=") for line in f.readlines())
diff --git a/lib/spack/spack/test/conftest.py b/lib/spack/spack/test/conftest.py
index f0e5f7fdb2..c97a830bbe 100644
--- a/lib/spack/spack/test/conftest.py
+++ b/lib/spack/spack/test/conftest.py
@@ -1920,3 +1920,21 @@ def _func(spec_str, tests=False):
         return concretized_specs_cache[key].copy()
 
     return _func
+
+
+@pytest.fixture
+def shell_as(shell):
+    if sys.platform != "win32":
+        yield
+        return
+    if shell not in ("pwsh", "bat"):
+        raise RuntimeError("Shell must be one of supported Windows shells (pwsh|bat)")
+    try:
+        # fetch and store old shell type
+        _shell = os.environ.get("SPACK_SHELL", None)
+        os.environ["SPACK_SHELL"] = shell
+        yield
+    finally:
+        # restore old shell if one was set
+        if _shell:
+            os.environ["SPACK_SHELL"] = _shell
diff --git a/lib/spack/spack/test/util/environment.py b/lib/spack/spack/test/util/environment.py
index 801c2d19f4..481a58db47 100644
--- a/lib/spack/spack/test/util/environment.py
+++ b/lib/spack/spack/test/util/environment.py
@@ -113,13 +113,16 @@ def test_path_put_first(prepare_environment_for_tests):
     assert envutil.get_path("TEST_ENV_VAR") == expected
 
 
-def test_dump_environment(prepare_environment_for_tests, tmpdir):
+@pytest.mark.parametrize("shell", ["pwsh", "bat"] if sys.platform == "win32" else ["bash"])
+def test_dump_environment(prepare_environment_for_tests, shell_as, shell, tmpdir):
     test_paths = "/a:/b/x:/b/c"
     os.environ["TEST_ENV_VAR"] = test_paths
     dumpfile_path = str(tmpdir.join("envdump.txt"))
     envutil.dump_environment(dumpfile_path)
     with open(dumpfile_path, "r") as dumpfile:
-        if sys.platform == "win32":
+        if shell == "pwsh":
+            assert "$Env:TEST_ENV_VAR={}\n".format(test_paths) in list(dumpfile)
+        elif shell == "bat":
             assert 'set "TEST_ENV_VAR={}"\n'.format(test_paths) in list(dumpfile)
         else:
             assert "TEST_ENV_VAR={0}; export TEST_ENV_VAR\n".format(test_paths) in list(dumpfile)
@@ -164,11 +167,14 @@ def test_escape_double_quotes_in_shell_modifications():
 
     to_validate.set("QUOTED_VAR", '"MY_VAL"')
 
-    cmds = to_validate.shell_modifications()
-
-    if sys.platform != "win32":
+    if sys.platform == "win32":
+        cmds = to_validate.shell_modifications(shell="bat")
+        assert r'set "VAR=$PATH;$ANOTHER_PATH"' in cmds
+        assert r'set "QUOTED_VAR="MY_VAL"' in cmds
+        cmds = to_validate.shell_modifications(shell="pwsh")
+        assert r"$Env:VAR=$PATH;$ANOTHER_PATH" in cmds
+        assert r'$Env:QUOTED_VAR="MY_VAL"' in cmds
+    else:
+        cmds = to_validate.shell_modifications()
         assert 'export VAR="$PATH:$ANOTHER_PATH"' in cmds
         assert r'export QUOTED_VAR="\"MY_VAL\""' in cmds
-    else:
-        assert "export VAR=$PATH;$ANOTHER_PATH" in cmds
-        assert r'export QUOTED_VAR="MY_VAL"' in cmds
diff --git a/lib/spack/spack/util/environment.py b/lib/spack/spack/util/environment.py
index 3c63543551..da6832e4dd 100644
--- a/lib/spack/spack/util/environment.py
+++ b/lib/spack/spack/util/environment.py
@@ -47,6 +47,7 @@
     "csh": "setenv {0} {1};\n",
     "fish": "set -gx {0} {1};\n",
     "bat": 'set "{0}={1}"\n',
+    "pwsh": "$Env:{0}={1}\n",
 }
 
 
@@ -55,6 +56,7 @@
     "csh": "unsetenv {0};\n",
     "fish": "set -e {0};\n",
     "bat": 'set "{0}="\n',
+    "pwsh": "Remove-Item Env:{0}\n",
 }
 
 
@@ -172,7 +174,9 @@ def path_put_first(var_name: str, directories: List[Path]):
 
 
 def _win_env_var_to_set_line(var: str, val: str) -> str:
-    return f'set "{var}={val}"'
+    is_pwsh = os.environ.get("SPACK_SHELL", None) == "pwsh"
+    env_set_phrase = f"$Env:{var}={val}" if is_pwsh else f'set "{var}={val}"'
+    return env_set_phrase
 
 
 def _nix_env_var_to_source_line(var: str, val: str) -> str:
@@ -693,7 +697,7 @@ def apply_modifications(self, env: Optional[MutableMapping[str, str]] = None):
 
     def shell_modifications(
         self,
-        shell: str = "sh",
+        shell: str = "sh" if sys.platform != "win32" else os.environ.get("SPACK_SHELL", "bat"),
         explicit: bool = False,
         env: Optional[MutableMapping[str, str]] = None,
     ) -> str:
diff --git a/share/spack/setup-env.ps1 b/share/spack/setup-env.ps1
index edf9613ec5..d3bed93d9b 100644
--- a/share/spack/setup-env.ps1
+++ b/share/spack/setup-env.ps1
@@ -46,6 +46,10 @@ if ($null -eq $Env:EDITOR)
     $Env:EDITOR = "notepad"
 }
 
+# Set spack shell so we can detect powershell context
+$Env:SPACK_SHELL="pwsh"
+
+doskey /exename=powershell.exe spack=$Env:SPACK_ROOT\bin\spack.ps1 $args
 
 Write-Output "*****************************************************************"
 Write-Output "**************** Spack Package Manager **************************"
diff --git a/share/spack/spack-completion.bash b/share/spack/spack-completion.bash
index abf41cb052..2e98698e0f 100755
--- a/share/spack/spack-completion.bash
+++ b/share/spack/spack-completion.bash
@@ -935,7 +935,7 @@ _spack_env() {
 _spack_env_activate() {
     if $list_options
     then
-        SPACK_COMPREPLY="-h --help --sh --csh --fish --bat -v --with-view -V --without-view -p --prompt --temp -d --dir"
+        SPACK_COMPREPLY="-h --help --sh --csh --fish --bat --pwsh -v --with-view -V --without-view -p --prompt --temp -d --dir"
     else
         _environments
     fi