Add support for zsh

Author: jwodder

CHANGELOG.md

@@ -5,6 +5,7 @@ v0.2.0 (in development)
 - [#2] When outputting for Bash, backslashes in strings are now escaped
 - Added version, author, etc. variables to the top of the file
 - Added a `--version` option
+- Added support for zsh
 
 v0.1.0 (2018-09-09)
 -------------------

README.rst

@@ -1,5 +1,6 @@
-Here we have yet another script for Git-aware customization of the bash command
-prompt.  Unlike all the other scripts, I wrote this one, so it's better.
+Here we have yet another script for Git-aware customization of the command
+prompt in Bash and zsh.  Unlike all the other scripts, I wrote this one, so
+it's better.
 
 .. image:: screenshot.png
 
@@ -10,14 +11,15 @@ Features:
 - automatically truncates the current directory path if it gets too long
 - shows the status of the current Git repository (see below)
 - thoroughly documented and easily customizable
+- supports both Bash and zsh
 
 
 Requirements
 ============
 
 - Python 3, version 3.5 or higher
 - Git, version 1.7.10 or higher
-- bash
+- bash or zsh
 
 
 Installation & Setup
@@ -26,12 +28,19 @@ Installation & Setup
 1. Save `ps1.py <ps1.py>`_ to your computer somewhere (I put my copy at
    ``~/share/ps1.py``)
 
-2. Add the following line to the end of your ``~/.bashrc``:
+2. If using Bash, add the following line to the end of your ``~/.bashrc``:
 
    .. code:: shell
 
         PROMPT_COMMAND="$PROMPT_COMMAND"'; PS1="$(/usr/bin/python3 ~/share/ps1.py "${PS1_GIT:-}")"'
 
+   If using zsh, add the following to the end of your ``~/.zshrc``:
+
+   .. code:: shell
+
+        precmd_ps1_py() { PS1="$(/usr/bin/python3 ~/share/ps1.py --zsh "${PS1_GIT:-}")" }
+        precmd_functions+=( precmd_ps1_py )
+
    Replace ``/usr/bin/python3`` with the path to your Python 3 interpreter, and
    replace ``~/share/ps1.py`` with the location you saved ``ps1.py`` at as
    appropriate.
@@ -42,7 +51,7 @@ Installation & Setup
 
 5. If the Git integration causes you trouble (either because something breaks
    or just because it's taking too long to run), it can be temporarily disabled
-   by running ``PS1_GIT=off`` in bash.
+   by running ``PS1_GIT=off`` on the command line.
 
 
 Usage
@@ -54,8 +63,8 @@ Usage
 
 ``ps1.py`` outputs a single line containing a stylized prompt string for the
 current directory.  By default, the stylization is in a format usable in Bash's
-``PS1`` variable, though the ``--ansi`` option can be supplied to print the
-ANSI sequences directly instead.
+``PS1`` variable, though the ``--ansi`` and ``--zsh`` option can be supplied to
+change the format.
 
 ``ps1.py`` takes a single optional argument.  If this argument is "off", then
 the Git integration is disabled.  If it is any other value or not specified,
@@ -66,6 +75,7 @@ Options
 
 --ansi         Format output for direct display
 --bash         Format output for use in Bash's ``PS1`` (default)
+--zsh          Format output for use in zsh's ``PS1``
 -V, --version  Display version information and exit
 -h, --help     Display usage information and exit
 

ps1.py

@@ -1,17 +1,19 @@
 #!/usr/bin/python3
 """
-Yet another bash prompt script
+Yet another bash/zsh prompt script
 
-Here we have yet another script for Git-aware customization of the bash command
-prompt.  Unlike all the other scripts, I wrote this one, so it's better.
+Here we have yet another script for Git-aware customization of the command
+prompt in Bash and zsh.  Unlike all the other scripts, I wrote this one, so
+it's better.
 
 Features:
 
 - lets you know if you have mail in ``$MAIL``
 - shows chroot and `virtualenv <https://virtualenv.pypa.io>`_ prompt prefixes
 - automatically truncates the current directory path if it gets too long
-- shows the status of the current Git repository (see below)
+- shows the status of the current Git repository
 - thoroughly documented and easily customizable
+- supports both Bash and zsh
 
 Visit <https://github.com/jwodder/ps1.py> for more information.
 
@@ -22,21 +24,30 @@
 1. Save this script to your computer somewhere (I put my copy at
    ``~/share/ps1.py``)
 
-2. Add the following line to the end of your ``~/.bashrc``::
+2. If using Bash, add the following line to the end of your ``~/.bashrc``:
+
+   .. code:: shell
 
         PROMPT_COMMAND="$PROMPT_COMMAND"'; PS1="$(/usr/bin/python3 ~/share/ps1.py "${PS1_GIT:-}")"'
 
-    Replace ``/usr/bin/python3`` with the path to your Python 3 interpreter,
-    and replace ``~/share/ps1.py`` with the location you saved ``ps1.py`` at as
-    appropriate.
+   If using zsh, add the following to the end of your ``~/.zshrc``:
+
+   .. code:: shell
+
+        precmd_ps1_py() { PS1="$(/usr/bin/python3 ~/share/ps1.py --zsh "${PS1_GIT:-}")" }
+        precmd_functions+=( precmd_ps1_py )
+
+   Replace ``/usr/bin/python3`` with the path to your Python 3 interpreter, and
+   replace ``~/share/ps1.py`` with the location you saved ``ps1.py`` at as
+   appropriate.
 
 3. Open a new shell
 
 4. Enjoy!
 
 5. If the Git integration causes you trouble (either because something breaks
    or just because it's taking too long to run), it can be temporarily disabled
-   by running ``PS1_GIT=off`` in bash.
+   by running ``PS1_GIT=off`` on the command line.
 """
 
 __version__      = '0.2.0.dev1'
@@ -59,22 +70,30 @@
 
 class Color(Enum):
     """
-    An enum of the supported foreground colors.  Each color's value equals the
-    ANSI SGR parameter for setting that color as the foreground color.
+    An enum of the supported foreground colors.  Each color's value equals its
+    xterm number.
     """
 
-    RED           = 31
-    GREEN         = 32
-    YELLOW        = 33
-    BLUE          = 34
-    MAGENTA       = 35
-    CYAN          = 36
-    LIGHT_RED     = 91
-    LIGHT_GREEN   = 92
-    LIGHT_YELLOW  = 93
-    LIGHT_BLUE    = 94
-    LIGHT_MAGENTA = 95
-    LIGHT_CYAN    = 96
+    RED           = 1
+    GREEN         = 2
+    YELLOW        = 3
+    BLUE          = 4
+    MAGENTA       = 5
+    CYAN          = 6
+    LIGHT_RED     = 9
+    LIGHT_GREEN   = 10
+    LIGHT_YELLOW  = 11
+    LIGHT_BLUE    = 12
+    LIGHT_MAGENTA = 13
+    LIGHT_CYAN    = 14
+
+    def asfg(self):
+        """
+        Return the ANSI SGR parameter for setting the color as the foreground
+        color
+        """
+        c = self.value
+        return c + 30 if c < 8 else c + 82
 
 
 class BashStyler:
@@ -100,7 +119,7 @@ def __call__(self, s, fg=None, bold=False):
         s = self.escape(s)
         if fg is not None:
             s = r'\[\e[{}{}m\]{}\[\e[0m\]'.format(
-                fg.value,
+                fg.asfg(),
                 ';1' if bold else '',
                 s,
             )
@@ -134,15 +153,45 @@ def __call__(self, s, fg=None, bold=False):
         :param bool bold: whether to stylize the string as bold
         """
         if fg is not None:
-            s = '\033[{}{}m{}\033[0m'.format(fg.value, ';1' if bold else '', s)
+            s = '\033[{}{}m{}\033[0m'.format(fg.asfg(), ';1' if bold else '', s)
         elif bold:
             s = '\033[1m{}\033[0m'.format(s)
         return s
 
 
+class ZshStyler:
+    """ Class for escaping & styling strings for use in zsh's PS1 variable """
+
+    #: The actual prompt symbol to add at the end of the output, just before a
+    #: final space character
+    prompt_suffix = '%#'
+
+    def __call__(self, s, fg=None, bold=False):
+        """
+        Return the string ``s`` escaped for use in a zsh PS1 variable.  If
+        ``fg`` is non-`None`, the string will be wrapped in the proper escape
+        sequences to display it as the given foreground color.  If ``bold`` is
+        true, the string will be wrapped in the proper escape sequences to
+        display it bold.
+
+        :param str s: the string to stylize
+        :param Color fg: the foreground color to stylize the string with
+        :param bool bold: whether to stylize the string as bold
+        """
+        s = self.escape(s)
+        if bold:
+            s = '%B{}%b'.format(s)
+        if fg is not None:
+            s = '%F{{{}}}{}%f'.format(fg.value, s)
+        return s
+
+    def escape(self, s):
+        return s.replace('%', '%%')
+
+
 def main():
     parser = argparse.ArgumentParser(
-        description='Yet another bash prompt script.  '
+        description='Yet another bash/zsh prompt script.  '
                     f'Visit <{__url__}> for more information.'
     )
     parser.add_argument(
@@ -159,6 +208,13 @@ def main():
         const  = BashStyler,
         help   = "Format prompt for Bash's PS1 (default)",
     )
+    parser.add_argument(
+        '--zsh',
+        action = 'store_const',
+        dest   = 'stylecls',
+        const  = ZshStyler,
+        help   = "Format prompt for zsh's PS1",
+    )
     parser.add_argument(
         '-V', '--version',
         action  = 'version',