Do not color output to non-terminal; improve print() arg handling

diff --git a/src/text_helpers.py b/src/text_helpers.py
index 206e016..712609b 100644 (file)
--- a/src/text_helpers.py
+++ b/src/text_helpers.py
@@ -30,10 +30,11 @@ from functools import partial
 from itertools import islice
 
 from type_helpers import isstr
+from sys import stdout
 
 
 def head_and_tail(iterable, n_head=20, n_tail=20, n_elems=None,
-              skip_elem="...(skipping {n_skipped} elements)..."):
+                  skip_elem="...(skipping {n_skipped} elements)..."):
     """ convenient way to shorten a possibly very long iterable before printing
 
     Will not modify short iterables, but for longer lists/tuples/... will only
@@ -70,7 +71,7 @@ def head_and_tail(iterable, n_head=20, n_tail=20, n_elems=None,
     # yield skip element
     if n_elems > n_head + n_tail:
         if skip_elem is not None:
-            if (isstr(skip_elem)):
+            if isstr(skip_elem):
                 yield skip_elem.format(n_skipped=n_elems-n_head-n_tail)
             else:
                 yield skip_elem
@@ -117,6 +118,8 @@ _COLOR_TO_CODE = dict(zip([COLOR_BLACK, COLOR_RED, COLOR_GREEN, COLOR_YELLOW, \
 
 _ANSI_ESCAPE_SURROUND = '\x1b[{}m{}\x1b[0m'
 
+_STDOUT_IS_TTY = stdout.isatty()
+
 
 def colored(text, foreground=None, background=None, style=None):
     """ return text with given foreground/background ANSI color escape seqence
@@ -150,18 +153,54 @@ def colored(text, foreground=None, background=None, style=None):
     return _ANSI_ESCAPE_SURROUND.format(';'.join(prefixes), text)
 
 
-def print(*args, **kwargs):
-    _builtin_print(colored(*args, **kwargs))
+def print(text, *args, **kwargs):           # pylint: disable=redefined-builtin
+    """ like the builtin print function but also accepts color args
+
+    If any arg of :py:func:`colored` is given in `kwargs`, will run text with
+    color-args through that function. Runs built-in :py:builtin:`print`
+    function with result and other args.
+
+    ...todo:: color all args, not just first
+    """
+    foreground = None
+    background = None
+    style = None
+
+    # remove color info from kwargs
+    try:
+        foreground = kwargs['foreground']
+        del kwargs['foreground']
+    except KeyError:
+        pass
+
+    try:
+        background = kwargs['background']
+        del kwargs['background']
+    except KeyError:
+        pass
+
+    try:
+        style = kwargs['style']
+        del kwargs['style']
+    except KeyError:
+        pass
+
+    if _STDOUT_IS_TTY:
+        text_c = colored(text, foreground, background, style)
+    else:
+        text_c = text
+
+    _builtin_print(text_c, *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)
+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)
+cyan    = partial(print, foreground=COLOR_CYAN)
+white   = partial(print, foreground=COLOR_WHITE)
 
 normal    = partial(print, style=STYLE_NORMAL)
 bold      = partial(print, style=STYLE_BOLD)

diff --git a/test/test_text_helpers.py b/test/test_text_helpers.py
index a40c55d..5bf81e3 100644 (file)
--- a/test/test_text_helpers.py
+++ b/test/test_text_helpers.py
@@ -93,6 +93,7 @@ class TextHelpersTester(unittest.TestCase):
         expect = tuple(range(20)) + tuple(range(80,100))
         self.assertEqual(short, expect)
 
+    # todo: test that no color codes are printed if stdout is not a terminal
 
 if __name__ == '__main__':
     unittest.main()