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
This commit is contained in:
Todd Gamblin
2023-12-28 15:36:30 -08:00
committed by GitHub
parent fc1e0178bf
commit 379eeda576
2 changed files with 37 additions and 15 deletions

View File

@@ -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