zhangyiss/spack
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Moving utilities to a common LLNL package.

Browse Source
This commit is contained in:
Todd Gamblin
2014-03-12 22:24:47 -04:00
parent 03ee31e0e8
commit 9d01df9e8a
51 changed files with 229 additions and 227 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
lib/spack/llnl/__init__.py Normal file
View File

0
lib/spack/llnl/util/__init__.py Normal file
View File

0
lib/spack/llnl/util/compare/__init__.py Normal file
View File

70
lib/spack/llnl/util/compare/none_high.py Normal file
View File

@@ -0,0 +1,70 @@
##############################################################################
# Copyright (c) 2013, Lawrence Livermore National Security, LLC.
# Produced at the Lawrence Livermore National Laboratory.
#
# This file is part of Spack.
# Written by Todd Gamblin, tgamblin@llnl.gov, All rights reserved.
# LLNL-CODE-647188
#
# For details, see https://scalability-llnl.github.io/spack
# Please also see the LICENSE file for our notice and the LGPL.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License (as published by
# the Free Software Foundation) version 2.1 dated February 1999.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the IMPLIED WARRANTY OF
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the terms and
# conditions of the GNU General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
##############################################################################
"""
Functions for comparing values that may potentially be None.
These none_high functions consider None as greater than all other values.
"""
# Preserve builtin min and max functions
_builtin_min = min
_builtin_max = max
def lt(lhs, rhs):
"""Less-than comparison. None is greater than any value."""
return lhs != rhs and (rhs is None or (lhs is not None and lhs < rhs))
def le(lhs, rhs):
"""Less-than-or-equal comparison. None is greater than any value."""
return lhs == rhs or lt(lhs, rhs)
def gt(lhs, rhs):
"""Greater-than comparison. None is greater than any value."""
return lhs != rhs and not lt(lhs, rhs)
def ge(lhs, rhs):
"""Greater-than-or-equal comparison. None is greater than any value."""
return lhs == rhs or gt(lhs, rhs)
def min(lhs, rhs):
"""Minimum function where None is greater than any value."""
if lhs is None:
return rhs
elif rhs is None:
return lhs
else:
return _builtin_min(lhs, rhs)
def max(lhs, rhs):
"""Maximum function where None is greater than any value."""
if lhs is None or rhs is None:
return None
else:
return _builtin_max(lhs, rhs)

70
lib/spack/llnl/util/compare/none_low.py Normal file
View File

@@ -0,0 +1,70 @@
##############################################################################
# Copyright (c) 2013, Lawrence Livermore National Security, LLC.
# Produced at the Lawrence Livermore National Laboratory.
#
# This file is part of Spack.
# Written by Todd Gamblin, tgamblin@llnl.gov, All rights reserved.
# LLNL-CODE-647188
#
# For details, see https://scalability-llnl.github.io/spack
# Please also see the LICENSE file for our notice and the LGPL.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License (as published by
# the Free Software Foundation) version 2.1 dated February 1999.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the IMPLIED WARRANTY OF
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the terms and
# conditions of the GNU General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
##############################################################################
"""
Functions for comparing values that may potentially be None.
These none_low functions consider None as less than all other values.
"""
# Preserve builtin min and max functions
_builtin_min = min
_builtin_max = max
def lt(lhs, rhs):
"""Less-than comparison. None is lower than any value."""
return lhs != rhs and (lhs is None or (rhs is not None and lhs < rhs))
def le(lhs, rhs):
"""Less-than-or-equal comparison. None is less than any value."""
return lhs == rhs or lt(lhs, rhs)
def gt(lhs, rhs):
"""Greater-than comparison. None is less than any value."""
return lhs != rhs and not lt(lhs, rhs)
def ge(lhs, rhs):
"""Greater-than-or-equal comparison. None is less than any value."""
return lhs == rhs or gt(lhs, rhs)
def min(lhs, rhs):
"""Minimum function where None is less than any value."""
if lhs is None or rhs is None:
return None
else:
return _builtin_min(lhs, rhs)
def max(lhs, rhs):
"""Maximum function where None is less than any value."""
if lhs is None:
return rhs
elif rhs is None:
return lhs
else:
return _builtin_max(lhs, rhs)

90
lib/spack/llnl/util/filesystem.py Normal file
View File

@@ -0,0 +1,90 @@
##############################################################################
# Copyright (c) 2013, Lawrence Livermore National Security, LLC.
# Produced at the Lawrence Livermore National Laboratory.
#
# This file is part of Spack.
# Written by Todd Gamblin, tgamblin@llnl.gov, All rights reserved.
# LLNL-CODE-647188
#
# For details, see https://scalability-llnl.github.io/spack
# Please also see the LICENSE file for our notice and the LGPL.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License (as published by
# the Free Software Foundation) version 2.1 dated February 1999.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the IMPLIED WARRANTY OF
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the terms and
# conditions of the GNU General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
##############################################################################
import os
import re
import shutil
import errno
import getpass
from contextlib import contextmanager, closing
import llnl.util.tty as tty
from spack.util.compression import ALLOWED_ARCHIVE_TYPES
def install(src, dest):
"""Manually install a file to a particular location."""
tty.info("Installing %s to %s" % (src, dest))
shutil.copy(src, dest)
def expand_user(path):
"""Find instances of '%u' in a path and replace with the current user's
username."""
username = getpass.getuser()
if not username and '%u' in path:
tty.die("Couldn't get username to complete path '%s'" % path)
return path.replace('%u', username)
@contextmanager
def working_dir(dirname):
orig_dir = os.getcwd()
os.chdir(dirname)
yield
os.chdir(orig_dir)
def touch(path):
with closing(open(path, 'a')) as file:
os.utime(path, None)
def mkdirp(*paths):
for path in paths:
if not os.path.exists(path):
os.makedirs(path)
elif not os.path.isdir(path):
raise OSError(errno.EEXIST, "File alredy exists", path)
def join_path(prefix, *args):
path = str(prefix)
for elt in args:
path = os.path.join(path, str(elt))
return path
def ancestor(dir, n=1):
"""Get the nth ancestor of a directory."""
parent = os.path.abspath(dir)
for i in range(n):
parent = os.path.dirname(parent)
return parent
def can_access(file_name):
"""True if we have read/write access to the file."""
return os.access(file_name, os.R_OK|os.W_OK)

187
lib/spack/llnl/util/lang.py Normal file
View File

@@ -0,0 +1,187 @@
##############################################################################
# Copyright (c) 2013, Lawrence Livermore National Security, LLC.
# Produced at the Lawrence Livermore National Laboratory.
#
# This file is part of Spack.
# Written by Todd Gamblin, tgamblin@llnl.gov, All rights reserved.
# LLNL-CODE-647188
#
# For details, see https://scalability-llnl.github.io/spack
# Please also see the LICENSE file for our notice and the LGPL.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License (as published by
# the Free Software Foundation) version 2.1 dated February 1999.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the IMPLIED WARRANTY OF
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the terms and
# conditions of the GNU General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
##############################################################################
import os
import re
import sys
import functools
import inspect
# Ignore emacs backups when listing modules
ignore_modules = [r'^\.#', '~$']
def caller_locals():
"""This will return the locals of the *parent* of the caller.
This allows a fucntion to insert variables into its caller's
scope. Yes, this is some black magic, and yes it's useful
for implementing things like depends_on and provides.
"""
stack = inspect.stack()
try:
return stack[2][0].f_locals
finally:
del stack
def get_calling_package_name():
"""Make sure that the caller is a class definition, and return
the module's name. This is useful for getting the name of
spack packages from inside a relation function.
"""
stack = inspect.stack()
try:
# get calling function name (the relation)
relation = stack[1][3]
# Make sure locals contain __module__
caller_locals = stack[2][0].f_locals
finally:
del stack
if not '__module__' in caller_locals:
raise ScopeError(relation)
module_name = caller_locals['__module__']
base_name = module_name.split('.')[-1]
return base_name
def attr_required(obj, attr_name):
"""Ensure that a class has a required attribute."""
if not hasattr(obj, attr_name):
tty.die("No required attribute '%s' in class '%s'"
% (attr_name, obj.__class__.__name__))
def attr_setdefault(obj, name, value):
"""Like dict.setdefault, but for objects."""
if not hasattr(obj, name):
setattr(obj, name, value)
return getattr(obj, name)
def has_method(cls, name):
for base in inspect.getmro(cls):
if base is object:
continue
if name in base.__dict__:
return True
return False
def memoized(obj):
"""Decorator that caches the results of a function, storing them
in an attribute of that function."""
cache = obj.cache = {}
@functools.wraps(obj)
def memoizer(*args, **kwargs):
if args not in cache:
cache[args] = obj(*args, **kwargs)
return cache[args]
return memoizer
def list_modules(directory, **kwargs):
"""Lists all of the modules, excluding __init__.py, in a
particular directory. Listed packages have no particular
order."""
list_directories = kwargs.setdefault('directories', True)
for name in os.listdir(directory):
if name == '__init__.py':
continue
path = os.path.join(directory, name)
if list_directories and os.path.isdir(path):
init_py = os.path.join(path, '__init__.py')
if os.path.isfile(init_py):
yield name
elif name.endswith('.py'):
if not any(re.search(pattern, name) for pattern in ignore_modules):
yield re.sub('.py$', '', name)
def key_ordering(cls):
"""Decorates a class with extra methods that implement rich comparison
operations and __hash__. The decorator assumes that the class
implements a function called _cmp_key(). The rich comparison operations
will compare objects using this key, and the __hash__ function will
return the hash of this key.
If a class already has __eq__, __ne__, __lt__, __le__, __gt__, or __ge__
defined, this decorator will overwrite them. If the class does not
have a _cmp_key method, then this will raise a TypeError.
"""
def setter(name, value):
value.__name__ = name
setattr(cls, name, value)
if not has_method(cls, '_cmp_key'):
raise TypeError("'%s' doesn't define _cmp_key()." % cls.__name__)
setter('__eq__', lambda s,o: o is not None and s._cmp_key() == o._cmp_key())
setter('__lt__', lambda s,o: o is not None and s._cmp_key() < o._cmp_key())
setter('__le__', lambda s,o: o is not None and s._cmp_key() <= o._cmp_key())
setter('__ne__', lambda s,o: o is None or s._cmp_key() != o._cmp_key())
setter('__gt__', lambda s,o: o is None or s._cmp_key() > o._cmp_key())
setter('__ge__', lambda s,o: o is None or s._cmp_key() >= o._cmp_key())
setter('__hash__', lambda self: hash(self._cmp_key()))
return cls
@key_ordering
class HashableMap(dict):
"""This is a hashable, comparable dictionary. Hash is performed on
a tuple of the values in the dictionary."""
def _cmp_key(self):
return tuple(sorted(self.values()))
def copy(self):
"""Type-agnostic clone method. Preserves subclass type."""
# Construct a new dict of my type
T = type(self)
clone = T()
# Copy everything from this dict into it.
for key in self:
clone[key] = self[key].copy()
return clone
def in_function(function_name):
"""True if the caller was called from some function with
the supplied Name, False otherwise."""
stack = inspect.stack()
try:
for elt in stack[2:]:
if elt[3] == function_name:
return True
return False
finally:
del stack

95
lib/spack/llnl/util/tty/__init__.py Normal file
View File

@@ -0,0 +1,95 @@
##############################################################################
# Copyright (c) 2013, Lawrence Livermore National Security, LLC.
# Produced at the Lawrence Livermore National Laboratory.
#
# This file is part of Spack.
# Written by Todd Gamblin, tgamblin@llnl.gov, All rights reserved.
# LLNL-CODE-647188
#
# For details, see https://scalability-llnl.github.io/spack
# Please also see the LICENSE file for our notice and the LGPL.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License (as published by
# the Free Software Foundation) version 2.1 dated February 1999.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the IMPLIED WARRANTY OF
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the terms and
# conditions of the GNU General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
##############################################################################
import sys
from llnl.util.tty.color import *
debug = False
verbose = False
indent = " "
def msg(message, *args):
cprint("@*b{==>} %s" % cescape(message))
for arg in args:
print indent + str(arg)
def info(message, *args, **kwargs):
format = kwargs.get('format', '*b')
cprint("@%s{==>} %s" % (format, cescape(str(message))))
for arg in args:
print indent + str(arg)
def verbose(message, *args):
if verbose:
info(message, *args, format='c')
def debug(*args):
if debug:
info("Debug: " + message, *args, format='*g')
def error(message, *args):
info("Error: " + str(message), *args, format='*r')
def warn(message, *args):
info("Warning: " + str(message), *args, format='*Y')
def die(message, *args):
error(message, *args)
sys.exit(1)
def get_number(prompt, **kwargs):
default = kwargs.get('default', None)
abort = kwargs.get('abort', None)
if default is not None and abort is not None:
prompt += ' (default is %s, %s to abort) ' % (default, abort)
elif default is not None:
prompt += ' (default is %s) ' % default
elif abort is not None:
prompt += ' (%s to abort) ' % abort
number = None
while number is None:
ans = raw_input(prompt)
if ans == str(abort):
return None
if ans:
try:
number = int(ans)
if number < 1:
msg("Please enter a valid number.")
number = None
except ValueError:
msg("Please enter a valid number.")
elif default is not None:
number = default
return number

193
lib/spack/llnl/util/tty/colify.py Normal file
View File

@@ -0,0 +1,193 @@
##############################################################################
# Copyright (c) 2013, Lawrence Livermore National Security, LLC.
# Produced at the Lawrence Livermore National Laboratory.
#
# This file is part of Spack.
# Written by Todd Gamblin, tgamblin@llnl.gov, All rights reserved.
# LLNL-CODE-647188
#
# For details, see https://scalability-llnl.github.io/spack
# Please also see the LICENSE file for our notice and the LGPL.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License (as published by
# the Free Software Foundation) version 2.1 dated February 1999.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the IMPLIED WARRANTY OF
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the terms and
# conditions of the GNU General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
##############################################################################
# colify
# By Todd Gamblin, tgamblin@llnl.gov
#
# Takes a list of items as input and finds a good columnization of them,
# similar to how gnu ls does. You can pipe output to this script and
# get a tight display for it. This supports both uniform-width and
# variable-width (tighter) columns.
#
# Run colify -h for more information.
#
import os
import sys
import fcntl
import termios
import struct
def get_terminal_size():
"""Get the dimensions of the console."""
def ioctl_GWINSZ(fd):
try:
cr = struct.unpack('hh', fcntl.ioctl(fd, termios.TIOCGWINSZ, '1234'))
except:
return
return cr
cr = ioctl_GWINSZ(0) or ioctl_GWINSZ(1) or ioctl_GWINSZ(2)
if not cr:
try:
fd = os.open(os.ctermid(), os.O_RDONLY)
cr = ioctl_GWINSZ(fd)
os.close(fd)
except:
pass
if not cr:
cr = (os.environ.get('LINES', 25), os.environ.get('COLUMNS', 80))
return int(cr[1]), int(cr[0])
class ColumnConfig:
def __init__(self, cols):
self.cols = cols
self.line_length = 0
self.valid = True
self.widths = [0] * cols
def __repr__(self):
attrs = [(a,getattr(self, a)) for a in dir(self) if not a.startswith("__")]
return "<Config: %s>" % ", ".join("%s: %r" % a for a in attrs)
def config_variable_cols(elts, console_cols, padding):
# Get a bound on the most columns we could possibly have.
lengths = [len(elt) for elt in elts]
max_cols = max(1, console_cols / (min(lengths) + padding))
max_cols = min(len(elts), max_cols)
configs = [ColumnConfig(c) for c in xrange(1, max_cols+1)]
for elt, length in enumerate(lengths):
for i, conf in enumerate(configs):
if conf.valid:
col = elt / ((len(elts) + i) / (i + 1))
padded = length
if col < i:
padded += padding
if conf.widths[col] < padded:
conf.line_length += padded - conf.widths[col]
conf.widths[col] = padded
conf.valid = (conf.line_length < console_cols)
try:
config = next(conf for conf in reversed(configs) if conf.valid)
except StopIteration:
# If nothing was valid the screen was too narrow -- just use 1 col.
config = configs[0]
config.widths = [w for w in config.widths if w != 0]
config.cols = len(config.widths)
return config
def config_uniform_cols(elts, console_cols, padding):
max_len = max(len(elt) for elt in elts) + padding
cols = max(1, console_cols / max_len)
cols = min(len(elts), cols)
config = ColumnConfig(cols)
config.widths = [max_len] * cols
return config
def isatty(ostream):
force = os.environ.get('COLIFY_TTY', 'false').lower() != 'false'
return force or ostream.isatty()
def colify(elts, **options):
# Get keyword arguments or set defaults
output = options.get("output", sys.stdout)
indent = options.get("indent", 0)
padding = options.get("padding", 2)
# elts needs to be an array of strings so we can count the elements
elts = [str(elt) for elt in elts]
if not elts:
return
if not isatty(output):
for elt in elts:
output.write("%s\n" % elt)
return
console_cols = options.get("cols", None)
if not console_cols:
console_cols, console_rows = get_terminal_size()
elif type(console_cols) != int:
raise ValueError("Number of columns must be an int")
console_cols = max(1, console_cols - indent)
method = options.get("method", "variable")
if method == "variable":
config = config_variable_cols(elts, console_cols, padding)
elif method == "uniform":
config = config_uniform_cols(elts, console_cols, padding)
else:
raise ValueError("method must be one of: " + allowed_methods)
cols = config.cols
formats = ["%%-%ds" % width for width in config.widths[:-1]]
formats.append("%s") # last column has no trailing space
rows = (len(elts) + cols - 1) / cols
rows_last_col = len(elts) % rows
for row in xrange(rows):
output.write(" " * indent)
for col in xrange(cols):
elt = col * rows + row
output.write(formats[col] % elts[elt])
output.write("\n")
row += 1
if row == rows_last_col:
cols -= 1
if __name__ == "__main__":
import optparse
cols, rows = get_terminal_size()
parser = optparse.OptionParser()
parser.add_option("-u", "--uniform", action="store_true", default=False,
help="Use uniformly sized columns instead of variable-size.")
parser.add_option("-p", "--padding", metavar="PADDING", action="store",
type=int, default=2, help="Spaces to add between columns. Default is 2.")
parser.add_option("-i", "--indent", metavar="SPACES", action="store",
type=int, default=0, help="Indent the output by SPACES. Default is 0.")
parser.add_option("-w", "--width", metavar="COLS", action="store",
type=int, default=cols, help="Indent the output by SPACES. Default is 0.")
options, args = parser.parse_args()
method = "variable"
if options.uniform:
method = "uniform"
if sys.stdin.isatty():
parser.print_help()
sys.exit(1)
else:
colify([line.strip() for line in sys.stdin], method=method, **options.__dict__)

190
lib/spack/llnl/util/tty/color.py Normal file
View File

@@ -0,0 +1,190 @@
##############################################################################
# Copyright (c) 2013, Lawrence Livermore National Security, LLC.
# Produced at the Lawrence Livermore National Laboratory.
#
# This file is part of Spack.
# Written by Todd Gamblin, tgamblin@llnl.gov, All rights reserved.
# LLNL-CODE-647188
#
# For details, see https://scalability-llnl.github.io/spack
# Please also see the LICENSE file for our notice and the LGPL.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License (as published by
# the Free Software Foundation) version 2.1 dated February 1999.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the IMPLIED WARRANTY OF
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the terms and
# conditions of the GNU General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
##############################################################################
"""
This file implements an expression syntax, similar to printf, for adding
ANSI colors to text.
See colorize(), cwrite(), and cprint() for routines that can generate
colored output.
colorize will take a string and replace all color expressions with
ANSI control codes. If the isatty keyword arg is set to False, then
the color expressions will be converted to null strings, and the
returned string will have no color.
cwrite and cprint are equivalent to write() and print() calls in
python, but they colorize their output. If the stream argument is
not supplied, they write to sys.stdout.
Here are some example color expressions:
@r Turn on red coloring
@R Turn on bright red coloring
@*{foo} Bold foo, but don't change text color
@_{bar} Underline bar, but don't change text color
@*b Turn on bold, blue text
@_B Turn on bright blue text with an underline
@. Revert to plain formatting
@*g{green} Print out 'green' in bold, green text, then reset to plain.
@*ggreen@. Print out 'green' in bold, green text, then reset to plain.
The syntax consists of:
color-expr = '@' [style] color-code '{' text '}' | '@.' | '@@'
style = '*' | '_'
color-code = [krgybmcwKRGYBMCW]
text = .*
'@' indicates the start of a color expression. It can be followed
by an optional * or _ that indicates whether the font should be bold or
underlined. If * or _ is not provided, the text will be plain. Then
an optional color code is supplied. This can be [krgybmcw] or [KRGYBMCW],
where the letters map to black(k), red(r), green(g), yellow(y), blue(b),
magenta(m), cyan(c), and white(w). Lowercase letters denote normal ANSI
colors and capital letters denote bright ANSI colors.
Finally, the color expression can be followed by text enclosed in {}. If
braces are present, only the text in braces is colored. If the braces are
NOT present, then just the control codes to enable the color will be output.
The console can be reset later to plain text with '@.'.
To output an @, use '@@'. To output a } inside braces, use '}}'.
"""
import re
import sys
class ColorParseError(Exception):
"""Raised when a color format fails to parse."""
def __init__(self, message):
super(ColorParseError, self).__init__(message)
# Text styles for ansi codes
styles = {'*' : '1', # bold
'_' : '4', # underline
None : '0' } # plain
# Dim and bright ansi colors
colors = {'k' : 30, 'K' : 90, # black
'r' : 31, 'R' : 91, # red
'g' : 32, 'G' : 92, # green
'y' : 33, 'Y' : 93, # yellow
'b' : 34, 'B' : 94, # blue
'm' : 35, 'M' : 95, # magenta
'c' : 36, 'C' : 96, # cyan
'w' : 37, 'W' : 97 } # white
# Regex to be used for color formatting
color_re = r'@(?:@|\.|([*_])?([a-zA-Z])?(?:{((?:[^}]|}})*)})?)'
class match_to_ansi(object):
def __init__(self, color=True):
self.color = color
def escape(self, s):
"""Returns a TTY escape sequence for a color"""
if self.color:
return "\033[%sm" % s
else:
return ''
def __call__(self, match):
"""Convert a match object generated by color_re into an ansi color code
This can be used as a handler in re.sub.
"""
style, color, text = match.groups()
m = match.group(0)
if m == '@@':
return '@'
elif m == '@.':
return self.escape(0)
elif m == '@':
raise ColorParseError("Incomplete color format: '%s' in %s"
% (m, match.string))
string = styles[style]
if color:
if color not in colors:
raise ColorParseError("invalid color specifier: '%s' in '%s'"
% (color, match.string))
string += ';' + str(colors[color])
colored_text = ''
if text:
colored_text = text + self.escape(0)
return self.escape(string) + colored_text
def colorize(string, **kwargs):
"""Take a string and replace all color expressions with ANSI control
codes. Return the resulting string.
If color=False is supplied, output will be plain text without
control codes, for output to non-console devices.
"""
color = kwargs.get('color', True)
return re.sub(color_re, match_to_ansi(color), string)
def cwrite(string, stream=sys.stdout, color=None):
"""Replace all color expressions in string with ANSI control
codes and write the result to the stream. If color is
False, this will write plain text with o color. If True,
then it will always write colored output. If not supplied,
then it will be set based on stream.isatty().
"""
if color is None:
color = stream.isatty()
stream.write(colorize(string, color=color))
def cprint(string, stream=sys.stdout, color=None):
"""Same as cwrite, but writes a trailing newline to the stream."""
cwrite(string + "\n", stream, color)
def cescape(string):
"""Replace all @ with @@ in the string provided."""
return str(string).replace('@', '@@')
class ColorStream(object):
def __init__(self, stream, color=None):
self.__class__ = type(stream.__class__.__name__,
(self.__class__, stream.__class__), {})
self.__dict__ = stream.__dict__
self.color = color
self.stream = stream
def write(self, string, **kwargs):
if kwargs.get('raw', False):
super(ColorStream, self).write(string)
else:
cwrite(string, self.stream, self.color)
def writelines(self, sequence, **kwargs):
raw = kwargs.get('raw', False)
for string in sequence:
self.write(string, self.color, raw=raw)