shell: fix zsh color formatting for PS1 in environments (#39497)

* shell: fix zsh color formatting for PS1 in environments

The `colorize` function in `llnl.util.tty.color` only applies proper formatting for Bash
ANSI and for console output, but this is not what zsh expects for environment variables.

In particular, when using `zsh`, `spack env activate -p` produces a `PS1` prompt that
looks like this:

```
\[\033[0;92m\][ENVIRONMENT]\[\033[0m\]
```

For zsh the formatting should be:

```
\e[0;92m[ENVIRONMENT]\e0;m
```

- [x] Add a `zsh` option to `colorize()` to enable zsh color formatting
- [x] Add conditional to choose the right `PS1` for `zsh`, `bash`, and `sh`
- [x] Don't use color escapes for `sh`, as they don't print properly

* convert lots of += lines to triple quotes

2 files changed, 37 insertions, 15 deletions

diff --git a/lib/spack/llnl/util/tty/color.py b/lib/spack/llnl/util/tty/color.py
index 67bd129eab..6a1d73cb16 100644
--- a/lib/spack/llnl/util/tty/color.py
+++ b/lib/spack/llnl/util/tty/color.py
@@ -204,17 +204,23 @@ def color_when(value):
 
 
 class match_to_ansi:
-    def __init__(self, color=True, enclose=False):
+    def __init__(self, color=True, enclose=False, zsh=False):
         self.color = _color_when_value(color)
         self.enclose = enclose
+        self.zsh = zsh
 
     def escape(self, s):
         """Returns a TTY escape sequence for a color"""
         if self.color:
-            if self.enclose:
-                return r"\[\033[%sm\]" % s
+            if self.zsh:
+                result = rf"\e[0;{s}m"
             else:
-                return "\033[%sm" % s
+                result = f"\033[{s}m"
+
+            if self.enclose:
+                result = rf"\[{result}\]"
+
+            return result
         else:
             return ""
 
@@ -261,9 +267,11 @@ def colorize(string, **kwargs):
             codes, for output to non-console devices.
         enclose (bool): If True, enclose ansi color sequences with
             square brackets to prevent misestimation of terminal width.
+        zsh (bool): If True, use zsh ansi codes instead of bash ones (for variables like PS1)
     """
     color = _color_when_value(kwargs.get("color", get_color_when()))
-    string = re.sub(color_re, match_to_ansi(color, kwargs.get("enclose")), string)
+    zsh = kwargs.get("zsh", False)
+    string = re.sub(color_re, match_to_ansi(color, kwargs.get("enclose")), string, zsh)
     string = string.replace("}}", "}")
     return string
 
diff --git a/lib/spack/spack/environment/shell.py b/lib/spack/spack/environment/shell.py
index a4f9634a8d..7ead2093a7 100644
--- a/lib/spack/spack/environment/shell.py
+++ b/lib/spack/spack/environment/shell.py
@@ -3,6 +3,7 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 import os
+import textwrap
 from typing import Optional
 
 import llnl.util.tty as tty
@@ -54,22 +55,35 @@ def activate_header(env, shell, prompt=None, view: Optional[str] = None):
         if view:
             cmds += "$Env:SPACK_ENV_VIEW='%s'\n" % view
     else:
-        if "color" in os.getenv("TERM", "") and prompt:
-            prompt = colorize("@G{%s}" % prompt, color=True, enclose=True)
+        bash_color_prompt = colorize(f"@G{{{prompt}}}", color=True, enclose=True)
+        zsh_color_prompt = colorize(f"@G{{{prompt}}}", color=True, enclose=False, zsh=True)
 
         cmds += "export SPACK_ENV=%s;\n" % env.path
         if view:
             cmds += "export SPACK_ENV_VIEW=%s;\n" % view
         cmds += "alias despacktivate='spack env deactivate';\n"
         if prompt:
-            cmds += "if [ -z ${SPACK_OLD_PS1+x} ]; then\n"
-            cmds += "    if [ -z ${PS1+x} ]; then\n"
-            cmds += "        PS1='$$$$';\n"
-            cmds += "    fi;\n"
-            cmds += '    export SPACK_OLD_PS1="${PS1}";\n'
-            cmds += "fi;\n"
-            cmds += 'export PS1="%s ${PS1}";\n' % prompt
-
+            cmds += textwrap.dedent(
+                rf"""
+                if [ -z ${{SPACK_OLD_PS1+x}} ]; then
+                    if [ -z ${{PS1+x}} ]; then
+                        PS1='$$$$';
+                    fi;
+                    export SPACK_OLD_PS1="${{PS1}}";
+                fi;
+                if [ -n "${{TERM:-}}" ] && [ "${{TERM#*color}}" != "${{TERM}}" ] && \
+                   [ -n "${{BASH:-}}" ];
+                then
+                    export PS1="{bash_color_prompt} ${{PS1}}";
+                elif [ -n "${{TERM:-}}" ] && [ "${{TERM#*color}}" != "${{TERM}}" ] && \
+                     [ -n "${{ZSH_NAME:-}}" ];
+                then
+                    export PS1="{zsh_color_prompt} ${{PS1}}";
+                else
+                    export PS1="{prompt} ${{PS1}}";
+                fi
+                """
+            ).lstrip("\n")
     return cmds