Create text_helpers with color/style functions for output

diff --git a/src/text_helpers.py b/src/text_helpers.py
new file mode 100644 (file)
index 0000000..3865934
--- /dev/null
+++ b/src/text_helpers.py
@@ -0,0 +1,100 @@
+#!/usr/bin/env python3
+
+""" Functions for printing colored text using ANSI escape sequences
+
+Not sure whether to save this as text_helpers.py or a shell_helpers.py or some
+other. Let's see what else could be added to any of those.
+
+Functions might cause trouble when combined, e.g.::
+
+    bold('this is bold and ' + red('red') + ' and ' + green('green'))
+
+will show the text "and green" not in bold. May have to try using specific
+end-of-color or end-of-style escape sequences instead of 0 (reset-everything).
+
+References:
+* http://stackoverflow.com/questions/287871/print-in-terminal-with-colors-using-python
+* https://en.wikipedia.org/wiki/ANSI_escape_code
+
+.. codeauthor:: Intra2net AG <info@intra2net>
+"""
+
+from __future__ import print_function
+
+from builtins import print as _builtin_print
+from functools import partial
+
+
+COLOR_BLACK = 'black'
+COLOR_RED = 'red'
+COLOR_GREEN = 'green'
+COLOR_YELLOW = 'yellow'
+COLOR_BLUE = 'blue'
+COLOR_MAGENTA = 'magenta'
+COLOR_CYAN = 'cyan'
+COLOR_WHITE = 'white'
+
+STYLE_NORMAL = 0
+STYLE_BOLD = 1
+STYLE_UNDERLINE = 4
+STYLE_BLINK = 5
+STYLE_REVERSE = 7
+
+
+_COLOR_TO_CODE = dict(zip([COLOR_BLACK, COLOR_RED, COLOR_GREEN, COLOR_YELLOW, \
+                           COLOR_BLUE, COLOR_MAGENTA, COLOR_CYAN, COLOR_WHITE],
+                          range(8)))
+
+_ANSI_ESCAPE_SURROUND = '\x1b[{}m{}\x1b[0m'
+
+
+def colored(text, foreground=None, background=None, style=None):
+    """ return text with given foreground/background ANSI color escape seqence
+
+    :param str text: text to color
+    :param str style: one of the style constants above
+    :param str foreground: one of the color constants to use for text color
+                           or None to leave as-is
+    :param str foreground: one of the color constants to use for background
+                           or None to leave as-is
+    :param style: single STYLE constant or iterable of those
+                  or None to leave as-is
+    :returns: text surrounded in ansi escape sequences
+    """
+
+    if foreground is None and background is None and style is None:
+        return text
+
+    prefixes = []
+    if foreground:
+        prefixes.append(str(30 + _COLOR_TO_CODE[foreground]))
+    if background:
+        prefixes.append(str(40 + _COLOR_TO_CODE[background]))
+    if style is None:
+        pass
+    elif isinstance(style, int):
+        prefixes.append(str(style))
+    else:   # assume iterable of ints
+        prefixes.extend(str(style_item) for style_item in style)
+
+    return _ANSI_ESCAPE_SURROUND.format(';'.join(prefixes), text)
+
+
+def print(*args, **kwargs):
+    _builtin_print(colored(*args, **kwargs))
+
+
+black = partial(print, foreground=COLOR_BLACK)
+red = partial(print, foreground=COLOR_RED)
+green = partial(print, foreground=COLOR_GREEN)
+yellow = partial(print, foreground=COLOR_YELLOW)
+blue = partial(print, foreground=COLOR_BLUE)
+magenta = partial(print, foreground=COLOR_MAGENTA)
+cyan = partial(print, foreground=COLOR_CYAN)
+white = partial(print, foreground=COLOR_WHITE)
+
+normal    = partial(print, style=STYLE_NORMAL)
+bold      = partial(print, style=STYLE_BOLD)
+underline = partial(print, style=STYLE_UNDERLINE)
+blink     = partial(print, style=STYLE_BLINK)
+reverse   = partial(print, style=STYLE_REVERSE)