bin/powerline-shell.py

651 lines
19 KiB
Python
Executable File
Raw Blame History

This file contains invisible Unicode characters

This file contains invisible Unicode characters that are indistinguishable to humans but may be processed differently by a computer. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import print_function
import argparse
import os
import sys
py3 = sys.version_info.major == 3
def warn(msg):
print('[powerline-bash] ', msg)
if py3:
def unicode(x):
return x
class Powerline:
symbols = {
'compatible': {
'lock': 'RO',
'network': 'SSH',
'separator': u'\u25B6',
'separator_thin': u'\u276F'
},
'patched': {
'lock': u'\uE0A2',
'network': u'\uE0A2',
'separator': u'\uE0B0',
'separator_thin': u'\uE0B1'
},
'flat': {
'lock': '',
'network': '',
'separator': '',
'separator_thin': ''
},
}
color_templates = {
'bash': '\\[\\e%s\\]',
'zsh': '%%{%s%%}',
'bare': '%s',
}
def __init__(self, args, cwd):
self.args = args
self.cwd = cwd
mode, shell = args.mode, args.shell
self.color_template = self.color_templates[shell]
self.reset = self.color_template % '[0m'
self.lock = Powerline.symbols[mode]['lock']
self.network = Powerline.symbols[mode]['network']
self.separator = Powerline.symbols[mode]['separator']
self.separator_thin = Powerline.symbols[mode]['separator_thin']
self.segments = []
def color(self, prefix, code):
if code is None:
return ''
elif code == Color.RESET:
return self.reset
else:
return self.color_template % ('[%s;5;%sm' % (prefix, code))
def fgcolor(self, code):
return self.color('38', code)
def bgcolor(self, code):
return self.color('48', code)
def append(self, content, fg, bg, separator=None, separator_fg=None):
self.segments.append((content, fg, bg,
separator if separator is not None else self.separator,
separator_fg if separator_fg is not None else bg))
def draw(self):
text = (''.join(self.draw_segment(i) for i in range(len(self.segments)))
+ self.reset) + ' '
if py3:
return text
else:
return text.encode('utf-8')
def draw_segment(self, idx):
segment = self.segments[idx]
next_segment = self.segments[idx + 1] if idx < len(self.segments)-1 else None
return ''.join((
self.fgcolor(segment[1]),
self.bgcolor(segment[2]),
segment[0],
self.bgcolor(next_segment[2]) if next_segment else self.reset,
self.fgcolor(segment[4]),
segment[3]))
class RepoStats:
symbols = {
'detached': u'\u2693',
'ahead': u'\u2B06',
'behind': u'\u2B07',
'staged': u'\u2714',
'not_staged': u'\u270E',
'untracked': u'\u2753',
'conflicted': u'\u273C'
}
def __init__(self):
self.ahead = 0
self.behind = 0
self.untracked = 0
self.not_staged = 0
self.staged = 0
self.conflicted = 0
@property
def dirty(self):
qualifiers = [
self.untracked,
self.not_staged,
self.staged,
self.conflicted,
]
return sum(qualifiers) > 0
def __getitem__(self, _key):
return getattr(self, _key)
def n_or_empty(self, _key):
"""Given a string name of one of the properties of this class, returns
the value of the property as a string when the value is greater than
1. When it is not greater than one, returns an empty string.
As an example, if you want to show an icon for untracked files, but you
only want a number to appear next to the icon when there are more than
one untracked files, you can do:
segment = repo_stats.n_or_empty("untracked") + icon_string
"""
return unicode(self[_key]) if int(self[_key]) > 1 else u''
def add_to_powerline(self, powerline, color):
def add(_key, fg, bg):
if self[_key]:
s = u" {}{} ".format(self.n_or_empty(_key), self.symbols[_key])
powerline.append(s, fg, bg)
add('ahead', color.GIT_AHEAD_FG, color.GIT_AHEAD_BG)
add('behind', color.GIT_BEHIND_FG, color.GIT_BEHIND_BG)
add('staged', color.GIT_STAGED_FG, color.GIT_STAGED_BG)
add('not_staged', color.GIT_NOTSTAGED_FG, color.GIT_NOTSTAGED_BG)
add('untracked', color.GIT_UNTRACKED_FG, color.GIT_UNTRACKED_BG)
add('conflicted', color.GIT_CONFLICTED_FG, color.GIT_CONFLICTED_BG)
def get_valid_cwd():
""" We check if the current working directory is valid or not. Typically
happens when you checkout a different branch on git that doesn't have
this directory.
We return the original cwd because the shell still considers that to be
the working directory, so returning our guess will confuse people
"""
# Prefer the PWD environment variable. Python's os.getcwd function follows
# symbolic links, which is undesirable. But if PWD is not set then fall
# back to this func
try:
cwd = os.getenv('PWD') or os.getcwd()
except:
warn("Your current directory is invalid. If you open a ticket at " +
"https://github.com/milkbikis/powerline-shell/issues/new " +
"we would love to help fix the issue.")
sys.stdout.write("> ")
sys.exit(1)
parts = cwd.split(os.sep)
up = cwd
while parts and not os.path.exists(up):
parts.pop()
up = os.sep.join(parts)
if cwd != up:
warn("Your current directory is invalid. Lowest valid directory: "
+ up)
return cwd
if __name__ == "__main__":
arg_parser = argparse.ArgumentParser()
arg_parser.add_argument('--cwd-mode', action='store',
help='How to display the current directory', default='fancy',
choices=['fancy', 'plain', 'dironly'])
arg_parser.add_argument('--cwd-only', action='store_true',
help='Deprecated. Use --cwd-mode=dironly')
arg_parser.add_argument('--cwd-max-depth', action='store', type=int,
default=5, help='Maximum number of directories to show in path')
arg_parser.add_argument('--cwd-max-dir-size', action='store', type=int,
help='Maximum number of letters displayed for each directory in the path')
arg_parser.add_argument('--colorize-hostname', action='store_true',
help='Colorize the hostname based on a hash of itself.')
arg_parser.add_argument('--mode', action='store', default='patched',
help='The characters used to make separators between segments',
choices=['patched', 'compatible', 'flat'])
arg_parser.add_argument('--shell', action='store', default='bash',
help='Set this to your shell type', choices=['bash', 'zsh', 'bare'])
arg_parser.add_argument('prev_error', nargs='?', type=int, default=0,
help='Error code returned by the last command')
args = arg_parser.parse_args()
powerline = Powerline(args, get_valid_cwd())
class DefaultColor:
"""
This class should have the default colors for every segment.
Please test every new segment with this theme first.
"""
# RESET is not a real color code. It is used as in indicator
# within the code that any foreground / background color should
# be cleared
RESET = -1
USERNAME_FG = 250
USERNAME_BG = 240
USERNAME_ROOT_BG = 124
HOSTNAME_FG = 250
HOSTNAME_BG = 238
HOME_SPECIAL_DISPLAY = True
HOME_BG = 31 # blueish
HOME_FG = 15 # white
PATH_BG = 237 # dark grey
PATH_FG = 250 # light grey
CWD_FG = 254 # nearly-white grey
SEPARATOR_FG = 244
READONLY_BG = 124
READONLY_FG = 254
SSH_BG = 166 # medium orange
SSH_FG = 254
REPO_CLEAN_BG = 148 # a light green color
REPO_CLEAN_FG = 0 # black
REPO_DIRTY_BG = 161 # pink/red
REPO_DIRTY_FG = 15 # white
JOBS_FG = 39
JOBS_BG = 238
CMD_PASSED_BG = 236
CMD_PASSED_FG = 15
CMD_FAILED_BG = 161
CMD_FAILED_FG = 15
SVN_CHANGES_BG = 148
SVN_CHANGES_FG = 22 # dark green
GIT_AHEAD_BG = 240
GIT_AHEAD_FG = 250
GIT_BEHIND_BG = 240
GIT_BEHIND_FG = 250
GIT_STAGED_BG = 22
GIT_STAGED_FG = 15
GIT_NOTSTAGED_BG = 130
GIT_NOTSTAGED_FG = 15
GIT_UNTRACKED_BG = 52
GIT_UNTRACKED_FG = 15
GIT_CONFLICTED_BG = 9
GIT_CONFLICTED_FG = 15
VIRTUAL_ENV_BG = 35 # a mid-tone green
VIRTUAL_ENV_FG = 00
class Color(DefaultColor):
"""
This subclass is required when the user chooses to use 'default' theme.
Because the segments require a 'Color' class for every theme.
"""
pass
class DefaultColor:
"""
This class should have the default colors for every segment.
Please test every new segment with this theme first.
"""
# RESET is not a real color code. It is used as in indicator
# within the code that any foreground / background color should
# be cleared
RESET = -1
USERNAME_FG = 223 #fg1
USERNAME_BG = 237 #bg1
USERNAME_ROOT_BG = 172 #yellow (dark)
HOSTNAME_FG = 223 #fg1
HOSTNAME_BG = 235 #bg0
HOME_SPECIAL_DISPLAY = False
# HOME_BG = 239 #bg2
HOME_FG = 223 #fg1
PATH_BG = 239 #bg2
PATH_FG = 223 #fg1
CWD_FG = 223 #fg1
SEPARATOR_FG = 246 #fg4
READONLY_BG = 124 #red (dark)
READONLY_FG = 223 #fg1
SSH_BG = 132 #purple (dark)
SSH_FG = 223 #fg1
REPO_CLEAN_BG = 106 #green (dark)
REPO_CLEAN_FG = 223 #fg1
REPO_DIRTY_BG = 124 #red (dark)
REPO_DIRTY_FG = 223 #fg1
JOBS_FG = 223 #fg1
JOBS_BG = 66 #blue (dark)
CMD_PASSED_BG = 142 #green (light)
CMD_PASSED_FG = 235 #bg0
CMD_FAILED_BG = 167 #red (light)
CMD_FAILED_FG = 235 #bg0
SVN_CHANGES_BG = 245 #gray
SVN_CHANGES_FG = 214 #yellow (light)
GIT_AHEAD_BG = 241 #bg3
GIT_AHEAD_FG = 223 #fg1
GIT_BEHIND_BG = 241 #bg3
GIT_BEHIND_FG = 223 #fg1
GIT_STAGED_BG = 142 #green (light)
GIT_STAGED_FG = 235 #bg0
GIT_NOTSTAGED_BG = 214 #yellow (light)
GIT_NOTSTAGED_FG = 235 #bg0
GIT_UNTRACKED_BG = 167 #red (light)
GIT_UNTRACKED_FG = 235 #bg0
GIT_CONFLICTED_BG = 167 #red (light)
GIT_CONFLICTED_FG = 124 #red (dark)
VIRTUAL_ENV_BG = 106 #green (dark)
VIRTUAL_ENV_FG = 223 #fg1
class Color(DefaultColor):
"""
This subclass is required when the user chooses to use 'default' theme.
Because the segments require a 'Color' class for every theme.
"""
pass
def add_set_term_title_segment(powerline):
term = os.getenv('TERM')
if not (('xterm' in term) or ('rxvt' in term)):
return
if powerline.args.shell == 'bash':
set_title = '\\[\\e]0;\\u@\\h: \\w\\a\\]'
elif powerline.args.shell == 'zsh':
set_title = '%{\033]0;%n@%m: %~\007%}'
else:
import socket
set_title = '\033]0;%s@%s: %s\007' % (os.getenv('USER'), socket.gethostname().split('.')[0], powerline.cwd or os.getenv('PWD'))
powerline.append(set_title, None, None, '')
add_set_term_title_segment(powerline)
def add_username_segment(powerline):
import os
if powerline.args.shell == 'bash':
user_prompt = ' \\u '
elif powerline.args.shell == 'zsh':
user_prompt = ' %n '
else:
user_prompt = ' %s ' % os.getenv('USER')
if os.getenv('USER') == 'root':
bgcolor = Color.USERNAME_ROOT_BG
else:
bgcolor = Color.USERNAME_BG
powerline.append(user_prompt, Color.USERNAME_FG, bgcolor)
add_username_segment(powerline)
import os
def add_ssh_segment(powerline):
if os.getenv('SSH_CLIENT'):
powerline.append(' %s ' % powerline.network, Color.SSH_FG, Color.SSH_BG)
add_ssh_segment(powerline)
import os
import re
import subprocess
import platform
def add_jobs_segment(powerline):
num_jobs = 0
if platform.system().startswith('CYGWIN'):
# cygwin ps is a special snowflake...
output_proc = subprocess.Popen(['ps', '-af'], stdout=subprocess.PIPE)
output = map(lambda l: int(l.split()[2].strip()),
output_proc.communicate()[0].decode("utf-8").splitlines()[1:])
num_jobs = output.count(os.getppid()) - 1
else:
pppid_proc = subprocess.Popen(['ps', '-p', str(os.getppid()), '-oppid='],
stdout=subprocess.PIPE)
pppid = pppid_proc.communicate()[0].decode("utf-8").strip()
output_proc = subprocess.Popen(['ps', '-a', '-o', 'ppid'],
stdout=subprocess.PIPE)
output = output_proc.communicate()[0].decode("utf-8")
num_jobs = len(re.findall(str(pppid), output)) - 1
if num_jobs > 0:
powerline.append(' %d ' % num_jobs, Color.JOBS_FG, Color.JOBS_BG)
add_jobs_segment(powerline)
import os
ELLIPSIS = u'\u2026'
def replace_home_dir(cwd):
home = os.getenv('HOME')
if cwd.startswith(home):
return '~' + cwd[len(home):]
return cwd
def split_path_into_names(cwd):
names = cwd.split(os.sep)
if names[0] == '':
names = names[1:]
if not names[0]:
return ['/']
return names
def requires_special_home_display(name):
"""Returns true if the given directory name matches the home indicator and
the chosen theme should use a special home indicator display."""
return (name == '~' and Color.HOME_SPECIAL_DISPLAY)
def maybe_shorten_name(powerline, name):
"""If the user has asked for each directory name to be shortened, will
return the name up to their specified length. Otherwise returns the full
name."""
if powerline.args.cwd_max_dir_size:
return name[:powerline.args.cwd_max_dir_size]
return name
def get_fg_bg(name, is_last_dir):
"""Returns the foreground and background color to use for the given name.
"""
if requires_special_home_display(name):
return (Color.HOME_FG, Color.HOME_BG,)
if is_last_dir:
return (Color.CWD_FG, Color.PATH_BG,)
else:
return (Color.PATH_FG, Color.PATH_BG,)
def add_cwd_segment(powerline):
cwd = powerline.cwd or os.getenv('PWD')
if not py3:
cwd = cwd.decode("utf-8")
cwd = replace_home_dir(cwd)
if powerline.args.cwd_mode == 'plain':
powerline.append(' %s ' % (cwd,), Color.CWD_FG, Color.PATH_BG)
return
names = split_path_into_names(cwd)
max_depth = powerline.args.cwd_max_depth
if max_depth <= 0:
warn("Ignoring --cwd-max-depth argument since it's not greater than 0")
elif len(names) > max_depth:
# https://github.com/milkbikis/powerline-shell/issues/148
# n_before is the number is the number of directories to put before the
# ellipsis. So if you are at ~/a/b/c/d/e and max depth is 4, it will
# show `~ a ... d e`.
#
# max_depth must be greater than n_before or else you end up repeating
# parts of the path with the way the splicing is written below.
n_before = 2 if max_depth > 2 else max_depth - 1
names = names[:n_before] + [ELLIPSIS] + names[n_before - max_depth:]
if (powerline.args.cwd_mode == 'dironly' or powerline.args.cwd_only):
# The user has indicated they only want the current directory to be
# displayed, so chop everything else off
names = names[-1:]
for i, name in enumerate(names):
is_last_dir = (i == len(names) - 1)
fg, bg = get_fg_bg(name, is_last_dir)
separator = powerline.separator_thin
separator_fg = Color.SEPARATOR_FG
if requires_special_home_display(name) or is_last_dir:
separator = None
separator_fg = None
powerline.append(' %s ' % maybe_shorten_name(powerline, name), fg, bg,
separator, separator_fg)
add_cwd_segment(powerline)
import re
import subprocess
import os
def get_PATH():
"""Normally gets the PATH from the OS. This function exists to enable
easily mocking the PATH in tests.
"""
return os.getenv("PATH")
def git_subprocess_env():
return {
# LANG is specified to ensure git always uses a language we are expecting.
# Otherwise we may be unable to parse the output.
"LANG": "C",
# https://github.com/milkbikis/powerline-shell/pull/126
"HOME": os.getenv("HOME"),
# https://github.com/milkbikis/powerline-shell/pull/153
"PATH": get_PATH(),
}
def parse_git_branch_info(status):
info = re.search('^## (?P<local>\S+?)''(\.{3}(?P<remote>\S+?)( \[(ahead (?P<ahead>\d+)(, )?)?(behind (?P<behind>\d+))?\])?)?$', status[0])
return info.groupdict() if info else None
def _get_git_detached_branch():
p = subprocess.Popen(['git', 'describe', '--tags', '--always'],
stdout=subprocess.PIPE, stderr=subprocess.PIPE,
env=git_subprocess_env())
detached_ref = p.communicate()[0].decode("utf-8").rstrip('\n')
if p.returncode == 0:
branch = u'{} {}'.format(RepoStats.symbols['detached'], detached_ref)
else:
branch = 'Big Bang'
return branch
def parse_git_stats(status):
stats = RepoStats()
for statusline in status[1:]:
code = statusline[:2]
if code == '??':
stats.untracked += 1
elif code in ('DD', 'AU', 'UD', 'UA', 'DU', 'AA', 'UU'):
stats.conflicted += 1
else:
if code[1] != ' ':
stats.not_staged += 1
if code[0] != ' ':
stats.staged += 1
return stats
def add_git_segment(powerline):
try:
p = subprocess.Popen(['git', 'status', '--porcelain', '-b'],
stdout=subprocess.PIPE, stderr=subprocess.PIPE,
env=git_subprocess_env())
except OSError:
# Popen will throw an OSError if git is not found
return
pdata = p.communicate()
if p.returncode != 0:
return
status = pdata[0].decode("utf-8").splitlines()
stats = parse_git_stats(status)
branch_info = parse_git_branch_info(status)
if branch_info:
stats.ahead = branch_info["ahead"]
stats.behind = branch_info["behind"]
branch = branch_info['local']
else:
branch = _get_git_detached_branch()
bg = Color.REPO_CLEAN_BG
fg = Color.REPO_CLEAN_FG
if stats.dirty:
bg = Color.REPO_DIRTY_BG
fg = Color.REPO_DIRTY_FG
powerline.append(' %s ' % branch, fg, bg)
stats.add_to_powerline(powerline, Color)
add_git_segment(powerline)
import os
def add_read_only_segment(powerline):
cwd = powerline.cwd or os.getenv('PWD')
if not os.access(cwd, os.W_OK):
powerline.append(' %s ' % powerline.lock, Color.READONLY_FG, Color.READONLY_BG)
add_read_only_segment(powerline)
def add_root_segment(powerline):
root_indicators = {
'bash': ' \\$ ',
'zsh': ' %# ',
'bare': ' $ ',
}
bg = Color.CMD_PASSED_BG
fg = Color.CMD_PASSED_FG
if powerline.args.prev_error != 0:
fg = Color.CMD_FAILED_FG
bg = Color.CMD_FAILED_BG
powerline.append(root_indicators[powerline.args.shell], fg, bg)
add_root_segment(powerline)
sys.stdout.write(powerline.draw())