# Author: Martin Matusiak <numerodix@gmail.com>
import difflib
import os
import re
import sys
import warnings
__all__ = [ # noqa
"black",
"blue",
"cyan",
"green",
"magenta",
"red",
"white",
"yellow",
"colorize",
"colorize_v2",
"wrap_string",
"get_code",
"get_code_v2",
"highlight_string",
"get_highlighter",
"strip_escapes",
"justify_formatted",
"colordiff",
"set_term_title",
"write_out",
"write_err",
"Colors",
]
# Don't write escapes to dumb terminals
_disabled = (not os.environ.get("TERM")) or (os.environ.get("TERM") == "dumb")
[docs]class Colors(object):
"""Container class for colors"""
[docs] @classmethod
def new(cls, colorname):
try:
cls._colorlist
except AttributeError:
cls._colorlist = []
newcls = type.__new__(type, colorname, (object,), {})
newcls.id = len(cls._colorlist)
cls._colorlist.append(newcls)
setattr(cls, colorname, newcls)
[docs] @classmethod
def iter(cls):
for color in cls._colorlist:
yield color
# Define Colors members
Colors.new("Black")
Colors.new("Red")
Colors.new("Green")
Colors.new("Yellow")
Colors.new("Blue")
Colors.new("Magenta")
Colors.new("Cyan")
Colors.new("White")
# Define coloring shorthands
def make_func(color):
def f(s, bold=False, reverse=False):
return colorize(s, color, bold=bold, reverse=reverse)
f.__doc__ = (
"""
Colorize string in %s
:param string s: The string to colorize.
:param bool bold: Whether to mark up in bold.
:param bool reverse: Whether to mark up in reverse video.
:rtype: string
"""
% color.__name__.lower()
)
return f
for color in Colors.iter():
globals()[color.__name__.lower()] = make_func(color)
# Define highlighting colors
highlights = [
Colors.Green,
Colors.Yellow,
Colors.Cyan,
Colors.Blue,
Colors.Magenta,
Colors.Red,
]
highlight_map = {}
for (n, h) in enumerate(highlights):
highlight_map[n] = [color for color in Colors.iter() if h == color].pop()
# Coloring functions
[docs]def get_highlighter(colorid):
"""
Map a color index to a highlighting color.
:param int colorid: The index.
:rtype: :class:`Colors`
"""
return highlight_map[colorid % len(highlights)]
[docs]def get_code(color, bold=False, reverse=False):
"""
Returns the escape code for styling with the given color,
in bold and/or reverse.
:param color: The color to use.
:type color: :class:`Colors` class
:param bool bold: Whether to mark up in bold.
:param bool reverse: Whether to mark up in reverse video.
:rtype: string
"""
if _disabled:
return ""
fmt = "0;0"
if bold and reverse:
fmt = "1;7"
elif reverse:
fmt = "0;7"
elif bold:
fmt = "0;1"
color = (color is not None) and ";3%s" % color.id or ""
return "\033[" + fmt + color + "m"
[docs]def get_code_v2(color, bold=False, reverse=False, underline=False, blink=False):
"""
Returns the escape code for styling with the given color,
in bold and/or reverse.
:param color: The color to use.
:type color: :class:`Colors` class
:param bool bold: Whether to mark up in bold.
:param bool underline: Whether to mark up in underline.
:param bool blink: Whether to mark up in blink.
:param bool reverse: Whether to mark up in reverse video.
:rtype: string
"""
if _disabled:
return ""
fmt = "0"
items = []
if bold:
items.append("1")
if underline:
items.append("4")
if blink:
items.append("5")
if reverse:
items.append("7")
if len(items) != 0:
fmt = ";".join(items)
color = (color is not None) and ";3%s" % color.id or ""
return "\033[" + fmt + color + "m"
[docs]def colorize(s, color, bold=False, reverse=False, start=None, end=None):
"""
Colorize a string with the color given.
:param string s: The string to colorize.
:param color: The color to use.
:type color: :class:`Colors` class
:param bool bold: Whether to mark up in bold.
:param bool reverse: Whether to mark up in reverse video.
:param int start: Index at which to start coloring.
:param int end: Index at which to end coloring.
:rtype: string
"""
start = start if start else 0
end = end if end else len(s)
before = s[:start]
between = s[start:end]
after = s[end:]
return "%s%s%s%s%s" % (
before,
get_code(color, bold=bold, reverse=reverse),
between,
get_code(None),
after,
)
[docs]def colorize_v2(
s,
color,
bold=False,
reverse=False,
underline=False,
blink=False,
start=None,
end=None,
):
"""
Colorize a string with the color given.
:param string s: The string to colorize.
:param color: The color to use.
:type color: :class:`Colors` class
:param bool bold: Whether to mark up in bold.
:param bool reverse: Whether to mark up in reverse video.
:param bool blink: Whether to mark up in blink.
:param bool reverse: Whether to mark up in reverse video.
:param int start: Index at which to start coloring.
:param int end: Index at which to end coloring.
:rtype: string
"""
start = start if start else 0
end = end if end else len(s)
before = s[:start]
between = s[start:end]
after = s[end:]
return "%s%s%s%s%s" % (
before,
get_code_v2(
color, bold=bold, underline=underline, blink=blink, reverse=reverse
),
between,
get_code_v2(None),
after,
)
[docs]def wrap_string(s, pos, color, bold=False, reverse=False):
"""
Colorize the string up to a position.
:param string s: The string to colorize.
:param int pos: The position at which to stop.
:param color: The color to use.
:type color: :class:`Colors` class
:param bool bold: Whether to mark up in bold.
:param bool reverse: Whether to mark up in reverse video.
:rtype: string
.. deprecated:: 0.2.2
This function has been deprecated in favor of :func:`colorize`.
"""
warnings.warn("wrap_string is deprecated", DeprecationWarning, 2)
if _disabled:
if pos == 0:
pos = 1
return s[: pos - 1] + "|" + s[pos:]
return "%s%s%s%s" % (
get_code(color, bold=bold, reverse=reverse),
s[:pos],
get_code(None),
s[pos:],
)
[docs]def highlight_string(s, *spanlists, **kw):
"""
Highlight spans in a string using a list of (begin, end) pairs. Each
list is treated as a layer of highlighting. Up to four layers of
highlighting are supported.
:param string s: The string to highlight
:param list spanlists: A list of tuples on the form ``[(begin, end)*]*``
:param kw: May include: `bold`, `reverse`, `color`, `colors` and `nocolor`
:rtype: string
.. deprecated:: 0.2.3
The `color` parameter has been deprecated in favor of `colors`.
"""
colors = kw.get("colors", [])
# pair span with color and id of the list -> (span, color, list_id)
tuples = []
for id, spanlist in enumerate(spanlists):
try:
color = colors[id]
except IndexError:
color = get_highlighter(id)
tuples.extend([(span, color, id) for span in spanlist])
# produce list of (pos,color,start_end,list_id) pairs
# (begin, Red, True, list_id) # start new color
# (end, Red, False, list_id) # end current color
markers = []
for i in tuples:
(begin, end), color, list_id = i
markers.append((begin, color, True, list_id))
markers.append((end, color, False, list_id))
def get_key(tup):
pos, color, start_end, list_id = tup
return pos
markers.sort(key=get_key)
# produce list of (pos, color, layer) pairs
codes = []
stack = []
for (pos, color, start_end, list_id) in markers:
# stack invariant : list_id1 < list_id2 => i1 < i2
if start_end:
inserted = False
for (i, (c, id)) in enumerate(stack):
if list_id < id:
stack.insert(i, (color, list_id))
inserted = True
break
if not inserted:
stack.append((color, list_id))
else:
stack.remove((color, list_id))
cur_color = None
if len(stack) > 0:
(cur_color, _) = stack[-1]
codes.append((pos, cur_color, len(stack)))
# apply codes to the string
cursor = 0
segments = []
for (pos, color, layer) in codes:
bold = False
reverse = False
# allow bold/reverse/nocolor styling as parameters
if color:
if kw.get("color"):
color = kw.get("color")
warnings.warn("color is deprecated", DeprecationWarning, 2)
elif kw.get("nocolor"):
color = None
bold = kw.get("bold") or bold
reverse = kw.get("reverse") or reverse
if layer == 2:
bold = True
if layer == 3:
reverse = True
if layer >= 4:
bold = True
reverse = True
segments.append(s[cursor:pos])
segments.append(get_code(color, bold=bold, reverse=reverse))
cursor = pos
segments.append(s[cursor:])
return "".join(segments)
[docs]def colordiff(x, y, color_x=Colors.Cyan, color_y=Colors.Green, debug=False):
"""
Formats a diff of two strings using the longest common subsequence by
highlighting characters that differ between the strings.
Returns the strings `x` and `y` with highlighting.
:param string x: The first string.
:param string y: The second string.
:param color_x: The color to use for the first string.
:type color_x: :class:`Colors` class
:param color_y: The color to use for the second string.
:type color_y: :class:`Colors` class
:param bool debug: Whether to print debug output underway.
:rtype: tuple
"""
def compute_seq(x, y):
"""SequenceMatcher computes the longest common contiguous subsequence
rather than the longest common subsequence, but this just causes the
diff to show more changed characters, the result is still correct"""
sm = difflib.SequenceMatcher(None, x, y)
seq = ""
for match in sm.get_matching_blocks():
seq += x[match.a : match.a + match.size]
return seq
def make_generator(it):
g = ((i, e) for (i, e) in enumerate(it))
def f():
try:
return next(g)
except StopIteration:
return (-1, None)
return f
def log(s):
if debug:
print(s) # pragma: no cover
seq = compute_seq(x, y)
log(">>> %s , %s -> %s" % (x, y, seq))
it_seq = make_generator(seq)
it_x = make_generator(x)
it_y = make_generator(y)
(sid, s) = it_seq()
(aid, a) = it_x()
(bid, b) = it_y()
x_spans = []
y_spans = []
while True:
if not any([s, a, b]):
break
# character the same in all sets
# -> unchanged
if s == a == b:
log(" %s" % s)
(sid, s) = it_seq()
(aid, a) = it_x()
(bid, b) = it_y()
# character the same in orig and common
# -> added in new
elif s == a:
log("+%s" % b)
y_spans.append((bid, bid + 1))
(bid, b) = it_y()
# character the same in new and common
# -> removed in orig
elif s == b:
log("-%s" % a)
x_spans.append((aid, aid + 1))
(aid, a) = it_x()
# character not the same (eg. case change)
# -> removed in orig and added in new
elif a != b:
if a:
log("-%s" % a)
x_spans.append((aid, aid + 1))
(aid, a) = it_x()
if b:
log("+%s" % b)
y_spans.append((bid, bid + 1))
(bid, b) = it_y()
x_fmt = highlight_string(x, x_spans, reverse=True, colors=[color_x])
y_fmt = highlight_string(y, y_spans, reverse=True, colors=[color_y])
return x_fmt, y_fmt
[docs]def strip_escapes(s):
"""
Strips escapes from the string.
:param string s: The string.
:rtype: string
"""
return re.sub("\033\[(?:(?:[0-9]*;)*)(?:[0-9]*m)", "", s)
# Output functions
[docs]def set_term_title(s):
"""
Set the title of a terminal window.
:param string s: The title.
"""
if not _disabled:
sys.stdout.write("\033]2;%s\007" % s)
def write_to(target, s):
# assuming we have escapes in the string
if not _disabled:
if not os.isatty(target.fileno()):
s = strip_escapes(s)
target.write(s)
target.flush()
[docs]def write_out(s):
"""
Write a string to ``sys.stdout``, strip escapes if output is a pipe.
:param string s: The title.
"""
write_to(sys.stdout, s)
[docs]def write_err(s):
"""
Write a string to ``sys.stderr``, strip escapes if output is a pipe.
:param string s: The title.
"""
write_to(sys.stderr, s)