Skip to content

Instantly share code, notes, and snippets.

@seanh

seanh/README.md

Last active Mar 27, 2021
Embed
What would you like to do?
My ZSH prompt in Python

My ZSH prompt in Python, with:

  • A horizontal rule (as wide as the terminal) before each prompt (when scrolling this makes it easier to see where one command's output ends and the next command begins)
  • Nicely truncated current working directory (like fish shell)
  • Active virtualenv
  • Git branch
  • Non-zero exit status
  • user@host only if connected over ssh

To use, put zshprompt.py on your $PATH and put source /path/to/zshprompt.zsh in your .zshrc file.

Requires Python 2.7 (so for example it'll work on Ubuntu 12.04 but not on 10.04) and psutil, on Ubuntu:

sudo apt-get install python-psutil

TODO: Get rid of psutil dependency (it's used to get the username). Support older versions of Python.

Screenshot

I would like to add a couple more features, including:

  • Allow the prompts to be customised with command-line format options
  • Fall back on a basic pure-zsh prompt if zshprompt.py exits with non-zero

Known issues:

  • Symlinks get expanded in the current working directory
  • Getting the git branch can be slow the first time you cd into a git repo
#!/usr/bin/env python2
"""Print out zsh prompts.
"""
import os
import os.path
import subprocess
import socket
import psutil
def _zero_width(s):
'''Return the given string, wrapped in zsh zero-width codes.
This tells zsh that the string is a zero-width string, eg. for prompt
alignment and cursor positioning purposes. For example, ANSI escape
sequences should be marked as zero-width.
'''
return "%{{{s}%}}".format(s=s)
def _foreground(s, color):
colors = {
'black': '\x1b[30m',
'red': '\x1b[31m',
'green': '\x1b[32m',
'yellow': '\x1b[33m',
'blue': '\x1b[34m',
'magenta': '\x1b[35m',
'cyan': '\x1b[36m',
'white': '\x1b[37m',
}
return "{color}{s}".format(color=_zero_width(colors[color]), s=s)
def _background(s, color):
colors = {
'black': '\x1b[40m',
'red': '\x1b[41m',
'green': '\x1b[42m',
'yellow': '\x1b[43m',
'blue': '\x1b[44m',
'magenta': '\x1b[45m',
'cyan': '\x1b[46m',
'white': '\x1b[47m',
}
return "{color}{s}".format(color=_zero_width(colors[color]), s=s)
def _bold(s):
return "{bold}{s}".format(bold=_zero_width("\x1b[1m"), s=s)
def _underline(s):
return "{underline}{s}".format(underline=_zero_width("\x1b[4m"), s=s)
def _reverse(s):
return "{reverse}{s}".format(reverse=_zero_width("\x1b[7m"), s=s)
def _reset(s):
return "{s}{reset}".format(s=s, reset=_zero_width("\x1b[0m"))
def color(s, foreground=None, background=None, bold=False, underline=False,
reverse=False):
'''Return the given string, wrapped in the given colour.
Foreground and background can be one of:
black, red, green, yellow, blue, magenta, cyan, white.
Also resets the colour and other attributes at the end of the string.
'''
if not s:
return s
if foreground:
s = _foreground(s, foreground)
if background:
s = _background(s, background)
if bold:
s = _bold(s)
if underline:
s = _underline(s)
if reverse:
s = _reverse(s)
s = _reset(s)
return s
def horizontal_rule(char='-'):
'''Return a long string of the given characters.
The string will be as long as the width of the user's terminal in
characters, and will have a newline at the end.
'''
width = os.popen('stty size', 'r').read().split()[1]
width = int(width)
return char * width + _zero_width('\n')
def shorten_path(path, max_length=20):
'''Return the given path, shortened if it's too long.
Parent directories will be collapsed, fish-style. Examples:
/home/seanh -> ~
/home/seanh/Projects/ckan/ckan/ckan -> ~/P/c/c/ckan
/home/seanh/Projects/ckan/ckan-> ~/Projects/ckan/ckan
'''
# Replace the user's homedir in path with ~
homedir = os.path.expanduser('~')
if path.startswith(homedir):
path = '~' + path[len(homedir):]
parts = path.split(os.sep)
# Remove empty strings.
parts = [part for part in parts if part]
path = os.sep.join(parts)
# Starting from the root dir, truncate each dir to just its first letter
# until the full path is < max_length or all the dirs have already been
# truncated. Never truncate the last dir.
while len(path) > max_length:
for i in range(0, len(parts) - 1):
part = parts[i]
if len(part) > 1:
part = part[0]
parts[i] = part
path = os.sep.join(parts)
continue
break
return path
def current_working_dir():
'''Return the full absolute path to the current working directory.'''
# Code for getting the current working directory, copied from
# <https://github.com/Lokaltog/powerline/>.
try:
try:
cwd = os.getcwdu()
except AttributeError:
cwd = os.getcwd()
except OSError as e:
if e.errno == 2:
# User most probably deleted the directory, this happens when
# removing files from Mercurial repos for example.
cwd = "[not found]"
else:
raise
return cwd
def virtualenv():
path = os.environ.get('VIRTUAL_ENV', '')
if path:
path = os.path.basename(path)
return path
def git_branch():
# Warning: subprocess.check_output() is new in Python 2.7.
try:
output = subprocess.check_output('git status'.split(),
stderr=subprocess.PIPE)
except subprocess.CalledProcessError:
# Non-zero return code, assume the current working dir is not in a git
# repo.
return ''
first_line = output.split('\n')[0]
branch_name = first_line.split(' ')[-1]
return branch_name
def user_name():
return psutil.Process(os.getpid()).username
def host_name():
return socket.gethostname()
def ssh_user_at_host():
if os.environ.get('SSH_CONNECTION'):
return "{user}@{host}".format(user=user_name(), host=host_name())
else:
return ''
def left_prompt():
'''Return my zsh left prompt.
'''
return "{hr}{cwd} ".format(
hr=color(horizontal_rule(' '), background='black'),
cwd=color(shorten_path(current_working_dir()), foreground='green')
)
def right_prompt(last_exit_status):
'''Return my zsh right prompt.
'''
if last_exit_status in (None, 0):
last_exit_status = ''
else:
last_exit_status = str(last_exit_status)
parts = [
color(last_exit_status, foreground='red'),
color(virtualenv(), foreground='blue'),
color(git_branch(), foreground='yellow'),
color(ssh_user_at_host(), foreground='white'),
]
# Remove empty strings from parts.
parts = [part for part in parts if part]
prompt = ' '.join(parts).strip()
return prompt
def main():
import argparse
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('side', metavar='left|right',
choices=('left', 'right'),
help="which zsh prompt to print (the left- or right-side prompt)")
parser.add_argument('--last-exit-status', dest='last_exit_status',
type=int,
help='the exit status (int) of the previous shell command '
'(default: None, printing last exit status will not be '
'supported)')
args = parser.parse_args()
if args.side == 'left':
print left_prompt()
else:
assert args.side == 'right'
print right_prompt(args.last_exit_status)
if __name__ == '__main__':
main()
PROMPT='$(zshprompt.py left)'
RPROMPT='$(zshprompt.py right --last-exit-status=$?)'
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment