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

1 files changed, 13 insertions, 5 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