Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
vcprompt
#!/usr/bin/env python
"""
Usage: vcprompt [options]
Version control information in your prompt.
Attribution: possible original author <Matthias Riegler https://github.com/xvzf>?
Options:
-f, --format FORMAT The format string to use.
-p, --path PATH The path to run vcprompt on.
-d, --max-depth DEPTH The maximum number of directories to traverse.
-u, --unknown UNKNOWN The "unknown" value.
-s, --systems Print all known VCSs and exit
-v, --version Show program's version number and exit
-h, --help Show this help message and exit
VCS-specific formatting:
These options can be used for VCS-specific prompt formatting.
--format-bzr FORMAT Bazaar
--format-cvs FORMAT CVS
--format-darcs FORMAT Darcs
--format-fossil FORMAT Fossil
--format-git FORMAT Git
--format-hg FORMAT Mercurial
--format-svn FORMAT Subversion
"""
from subprocess import call, Popen, PIPE
from xml.dom.minidom import parseString
import optparse
import os
import re
import sys
try:
import sqlite3
has_sqlite3 = True
except ImportError:
try:
from pysqlite2 import dbapi2 as sqlite3
except ImportError:
has_sqlite3 = False
__version__ = (0, 1, 5)
# check to make sure the '--without-environment' flag is called first
# this could be done in a callback, but we'd have to keep a note of every
# option which is affected by this flag
if '--without-environment' in sys.argv and \
sys.argv[1] != '--without-environment':
output = "The '--without-environment' option must come before any "
print >> sys.stderr, '%s other options.' % output
sys.exit(1)
# we need to get this in early because callbacks are always called after
# every other option is already set, regardless of when the callback option
# is actually used in the script
if len(sys.argv) > 1 and sys.argv[1] == '--without-environment':
for k in os.environ.keys():
if k.startswith('VCPROMPT'):
del os.environ[k]
del sys.argv[1]
# user editable options
DEPTH = os.environ.get('VCPROMPT_DEPTH', 0)
FORMAT = os.environ.get('VCPROMPT_FORMAT', '%s:%b')
UNKNOWN = os.environ.get('VCPROMPT_UNKNOWN', '(unknown)')
SYSTEMS = []
# status indicators
STAGED = '*'
MODIFIED = '+'
UNTRACKED = '?'
def helper(*args, **kwargs):
"""
Prints the module's docstring.
Doing this kills two birds with one stone: it adds PEP 257
compliance and allows us to stop using optparse's built-in
help flag.
"""
print __doc__.strip()
sys.exit(0)
def systems(*args, **kwargs):
"""Prints all available systems to stdout."""
for system in SYSTEMS:
print system.func_name
sys.exit(0)
def vcs(file):
"""
A convenience decorator.
Wraps a VCS function, appending it to ``SYSTEMS`` and sets the ``file``
attribute.
"""
def wrapper(function):
SYSTEMS.append(function)
function.file = file
return function
return wrapper
def version(*args):
"""
Convenience function for printing a version number.
"""
print 'vcprompt %s' % '.'.join(map(str, __version__))
sys.exit(0)
def vcprompt(options):
"""
Returns a formatted version control string for use in a shell prompt
or elsewhere.
Arguments:
``options``
An optparse.Values instance.
"""
options.path = os.path.abspath(os.path.expanduser(options.path))
prompt = None
count = 0
while options.path:
# bail out on non-existant paths
if not os.path.exists(options.path):
break
# We need to change the current working directory or the '--path'
# flag might not work correctly with some formatting args.
# It's easier to do this here, rather than in every VCS function
if options.path != os.getcwd():
os.chdir(options.path)
for vcs in SYSTEMS:
if not os.path.exists(vcs.file):
continue
# set up custom formatting
vcs_format = getattr(options, 'format-' + vcs.__name__, None)
if vcs_format:
options.format = vcs_format
# the "vcs" file
options.file = vcs.file
prompt = vcs(options)
if prompt is not None:
return prompt
if options.depth:
if count == options.depth:
break
count += 1
options.path = options.path.rsplit('/', 1)[0]
return ''
def main():
# parser
parser = optparse.OptionParser()
# dump the provided --help option
parser.remove_option('--help')
# our own --help flag
parser.add_option('-h', '--help', action='callback', callback=helper)
# format
parser.add_option('-f', '--format', dest='format', default=FORMAT)
# path
parser.add_option('-p', '--path', dest='path', default='.')
# max depth
parser.add_option('-d', '--max-depth', dest='depth', type='int',
default=DEPTH)
# systems
parser.add_option('-s', '--systems', action='callback', callback=systems)
# unknown
parser.add_option('-u', '--unknown', dest='unknown', default=UNKNOWN)
# version
parser.add_option('-v', '--version', action='callback', callback=version)
# vcs-specific formatting
for system in SYSTEMS:
default = 'VCPROMPT_FORMAT_%s' % system.__name__.upper()
default = os.environ.get(default, None)
dest = 'format-%s' % system.__name__
flag = '--%s' % dest
parser.add_option(flag, dest=dest, default=default)
options, args = parser.parse_args()
output = vcprompt(options)
return output
@vcs(file=".bzr/branch/last-revision")
def bzr(options):
"""
Bazaar
The Bazaar version control system
"""
branch = revision = sha = modified = untracked = options.unknown
# local revision or global sha
if re.search('%(r|h)', options.format):
try:
fh = open(options.file, 'r')
for line in fh:
line = line.strip()
revision, sha = line.split(' ', 1)
# compensate for empty Bazaar repositories
if sha == 'null:':
sha = unknown
else:
sha = sha.rsplit('-', 1)[-1][:7]
break
except IOError:
pass
# status (modified/untracked)
if re.search('%[mu]', options.format):
command = 'bzr status --short'
process = Popen(command.split(), stdout=PIPE)
output = process.communicate()[0].strip()
returncode = process.returncode
if returncode == 0:
if output == '':
modified = ''
untracked = ''
else:
for line in output.split('\n'):
if line.startswith('M'):
modified = MODIFIED
elif line.startswith('?'):
untracked = UNTRACKED
# formatting
output = options.format
output = output.replace('%b', os.path.basename(options.path))
output = output.replace('%h', sha)
output = output.replace('%r', revision)
output = output.replace('%m', modified)
output = output.replace('%u', untracked)
output = output.replace('%s', 'bzr')
output = output.replace('%n', 'bzr')
return output
@vcs(file="CVS")
def cvs(options):
"""
CVS
Concurrent Versions System.
"""
output = options.format
output = output.replace('%b', options.unknown)
output = output.replace('%h', options.unknown)
output = output.replace('%r', options.unknown)
output = output.replace('%m', options.unknown)
output = output.replace('%u', options.unknown)
output = output.replace('%s', 'cvs')
output = output.replace('%n', 'cvs')
return output
@vcs(file="_darcs/hashed_inventory")
def darcs(options):
"""
Darcs
Distributed. Interactive. Smart.
"""
branch = sha = modified = untracked = options.unknown
# sha
if re.search('%(h|r)', options.format):
command = 'darcs changes --last 1 --xml'
process = Popen(command.split(), stdout=PIPE, stderr=PIPE)
output = process.communicate()[0]
returncode = process.returncode
if returncode == 0:
dom = parseString(output)
patch = dom.getElementsByTagName("patch")[0].getAttribute("hash")
sha = patch.rsplit('-', 1)[-1].split('.')[0][:7]
# branch
# darcs doesn't have in-repo local branching (yet), so just use
# the directory name for now
# see also: http://bugs.darcs.net/issue555
branch = os.path.basename(options.path)
# modified
if re.search('%[mu]', options.format):
command = 'darcs whatsnew -l -s'
process = Popen(command.split(), stdout=PIPE, stderr=PIPE)
output = process.communicate()[0]
returncode = process.returncode
if returncode == 1:
modified = ''
untracked = ''
elif returncode == 0:
for line in output:
line = line.strip()
if line.startswith('M'):
modified = MODIFIED
elif line.startswith('a'):
untracked = UNTRACKED
# formatting
output = options.format
output = output.replace('%b', branch)
output = output.replace('%h', sha)
output = output.replace('%r', sha)
output = output.replace('%m', modified)
output = output.replace('%u', untracked)
output = output.replace('%s', 'darcs')
output = output.replace('%n', 'darcs')
return output
@vcs(file="_FOSSIL_")
def fossil(options):
"""
Fossil
The Fossil version control system.
"""
branch = sha = modified = untracked = options.unknown
# all this just to get the repository file :(
repository = None
if has_sqlite3:
try:
conn = None
try:
query = "SELECT value FROM vvar where name = 'repository'"
conn = sqlite3.connect(options.file)
c = conn.cursor()
c.execute(query)
repository = c.fetchone()[0]
except sqlite3.OperationalError:
pass
finally:
if conn:
conn.close()
# grab the sha from the repo
if repository is not None and has_sqlite3:
try:
conn = None
try:
query = "SELECT uuid from blob ORDER BY rid DESC LIMIT 1"
conn = sqlite3.connect(repository)
c = conn.cursor()
c.execute(query)
sha = c.fetchone()[0][:7]
except sqlite3.OperationalError:
pass
finally:
if conn:
conn.close()
# branch
if sha != options.unknown and has_sqlite3:
try:
conn = None
try:
query = """SELECT value FROM tagxref WHERE rid =
(SELECT rid FROM blob WHERE uuid LIKE '%s%%')
AND value is not NULL LIMIT 1 """ % sha
conn = sqlite3.connect(repository)
c = conn.cursor()
c.execute(query)
branch = c.fetchone()[0]
except (sqlite3.OperationalError, TypeError):
pass
finally:
if conn:
conn.close()
# modified
if '%m' in options.format:
command = 'fossil changes'
process = Popen(command.split(), stdout=PIPE)
output = process.communicate()[0]
returncode = process.returncode
if returncode == 0:
if output:
modified = MODIFIED
else:
modified = ''
# untracked files
if '%u' in options.format:
command = 'fossil extras'
process = Popen(command.split(), stdout=PIPE)
output = process.communicate()[0]
returncode = process.returncode
if returncode == 0:
if output:
untracked = UNTRACKED
else:
untracked = ''
# parse out formatting string
output = options.format
output = output.replace('%b', branch)
output = output.replace('%h', sha)
output = output.replace('%r', sha)
output = output.replace('%m', modified)
output = output.replace('%u', untracked)
output = output.replace('%s', 'fossil')
output = output.replace('%n', 'fossil')
return output
@vcs(file=".git")
def git(options):
"""
Git
The fast version control system.
"""
staged = branch = sha = modified = untracked = options.unknown
def revstring(ref, chars=7):
if not os.path.exists(ref):
return ''
try:
fh = open(ref, 'r')
for line in fh:
return line.strip()[0:chars]
except IOError:
pass
return ''
# the current branch is required to get the sha
if re.search('%(b|r|h)', options.format):
branch_file = os.path.join(options.file, 'HEAD')
try:
fh = open(branch_file, 'r')
for line in fh:
line = line.strip()
if line.startswith('ref: refs/heads/'):
branch = (line.split('/')[-1] or options.unknown)
break
except IOError:
pass
# sha/revision
if re.search('%(r|h)', options.format) and branch != options.unknown:
sha_file = os.path.join(options.file, 'refs/heads/%s' % branch)
sha = revstring(sha_file)
# modified
if '%m' in options.format:
command = 'git diff --name-status --exit-code'
returncode = call(command, stdout=PIPE, stderr=PIPE, shell=True)
if returncode == 1:
modified = MODIFIED
else:
modified = ''
# untracked files
if '%u' in options.format:
command = 'git ls-files --other --exclude-standard'
process = Popen(command.split(), stdout=PIPE, stderr=PIPE)
output = process.communicate()[0]
returncode = process.returncode
if returncode == 0:
if output == '':
untracked = ''
else:
untracked = UNTRACKED
# staged files
if '%a' in options.format:
command = 'git diff --name-only --cached'
process = Popen(command.split(), stdout=PIPE, stderr=PIPE)
output = process.communicate()[0]
returncode = process.returncode
if returncode == 0:
if output == '':
staged = ''
else:
staged = STAGED
# formatting
output = options.format
output = output.replace('%b', branch)
output = output.replace('%h', sha)
output = output.replace('%r', sha)
output = output.replace('%m', modified)
output = output.replace('%u', untracked)
output = output.replace('%a', staged)
output = output.replace('%s', 'git')
output = output.replace('%n', 'git')
return output
@vcs(file=".hg")
def hg(options):
"""
Mercurial
The Mercurial version control system.
"""
branch = revision = sha = modified = untracked = options.unknown
# changeset ID or global sha
if re.search('%(r|h)', options.format):
try:
fh = open(os.path.join(options.file, 'cache/branchheads'), 'r')
line = fh.readline()
sha, revision = line.strip().split()
sha = sha[:7]
except IOError:
pass
# branch
if '%b' in options.format:
file = os.path.join(options.file, 'undo.branch')
try:
fh = open(file, 'r')
line = fh.readline()
branch = line.strip()
except IOError:
pass
# modified
if '%m' in options.format:
command = 'hg status --modified'
process = Popen(command.split(), stdout=PIPE, stderr=PIPE)
output = process.communicate()[0].strip()
returncode = process.returncode
if returncode == 0:
if output == '':
modified = ''
else:
modified = MODIFIED
# untracked
if '%u' in options.format:
command = 'hg status --unknown'
process = Popen(command.split(), stdout=PIPE, stderr=PIPE)
output = process.communicate()[0]
returncode = process.returncode
if output == '':
untracked = ''
else:
untracked = UNTRACKED
output = options.format
output = output.replace('%b', branch)
output = output.replace('%h', sha)
output = output.replace('%r', revision)
output = output.replace('%m', modified)
output = output.replace('%u', untracked)
output = output.replace('%s', 'hg')
output = output.replace('%n', 'hg')
return output
@vcs(file=".svn/entries")
def svn(options):
"""
Subversion
The Subversion version control system.
"""
branch = revision = modified = untracked = options.unknown
# branch
command = 'svn info %s' % options.path
process = Popen(command.split(), stdout=PIPE, stderr=PIPE)
output = process.communicate()[0]
returncode = process.returncode
if returncode == 0:
# compile some regexes
branch_regex = re.compile('((tags|branches)|trunk)')
revision_regex = re.compile('^Revision: (?P<revision>\d+)')
for line in output.split('\n'):
# branch
if '%b' in options.format:
if re.match('URL:', line):
matches = re.search(branch_regex, line)
if matches:
branch = matches.groups(0)[0]
# revision/sha
if re.search('%(r|h)', options.format):
if re.match('Revision:', line):
matches = re.search(revision_regex, line)
if 'revision' in matches.groupdict():
revision = matches.group('revision')
# modified
if re.search('%[mu]', options.format):
command = 'svn status'
process = Popen(command, shell=True, stdout=PIPE)
output = process.communicate()[0]
returncode = process.returncode
if returncode == 0:
if not output:
modified = ''
untracked = ''
else:
codes = [line[0] for line in output.split('\n') if line]
if 'M' in codes:
modified = MODIFIED
if '?' in codes:
untracked = UNTRACKED
# formatting
output = options.format
output = output.replace('%r', revision)
output = output.replace('%h', revision)
output = output.replace('%b', branch)
output = output.replace('%m', modified)
output = output.replace('%u', untracked)
output = output.replace('%s', 'svn')
output = output.replace('%n', 'svn')
return output
if __name__ == '__main__':
prompt = main()
if prompt:
print prompt
else:
sys.exit(1)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment