diff options
| -rw-r--r-- | lib/spack/llnl/util/tty/color.py | 18 | ||||
| -rw-r--r-- | lib/spack/spack/environment/shell.py | 34 |
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 |