Source code for ansicolor.ansicolor

# 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 justify_formatted(s, justify_func, width): """ Justify a formatted string to a width using a function (eg. ``string.ljust``). :param string s: The formatted string. :param justify_func: The justify function. :param int width: The width at which to justify. :rtype: string """ dx = len(s) - len(strip_escapes(s)) return justify_func(s, width + dx)
[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)