(python) colour printing with decorator in a function

Question

How can I decorate a function so that anything it prints to stdout is in green and anything it prints to stderr is in red? I have the termcolor module available.

Bonus karma: How can I pass arguments to the decorator to specify the colours, defaulting them to red and green?

Answer 1

An interesting question. The simplest solution would be similar to what Pete suggest. Just print the escape codes before running the function to each of stderr and stdout. However, if both stderr and stdout feed the same terminal, as is usually the case, they will interfere.

So, an alternative solution is to monkey-patch stdout and stderr with a tiny wrapper that enables the color for the duration of each write, taking care to do this only if we're in a terminal (and not piped).

#!/usr/bin/python2

import sys

def colorize(stdoutColor, stderrColor):
  defaultColor = '\033[0;0m'

  def applyColorize(f):
    class colorWrapper(object):
      def __init__(self, wrapee, color):
        self.wrapee = wrapee
        self.color = color
      def __getattr__(self, attr):
        if attr == 'write' and self.wrapee.isatty():
          return lambda x: self.wrapee.write(self.color + x + defaultColor)
        else:
          return getattr(self.wrapee, attr)

    def wrapper(*args, **kwds):
      oldStdout = sys.stdout
      oldStderr = sys.stderr
      sys.stdout = colorWrapper(oldStdout, stdoutColor)
      sys.stderr = colorWrapper(oldStderr, stderrColor)
      try:
        f(*args, **kwds)
      finally:
        sys.stdout = oldStdout
        sys.stderr = oldStderr

    return wrapper

  return applyColorize


greenColor = '\033[01;32m'
redColor = '\033[01;31m'

def foo():
  print "I'm ordinary and boring!"
  print >> sys.stderr, 'Writing to stderr!'

@colorize(greenColor, redColor)
def colorFoo():
  print "I'm colorful and exciting!"
  print >> sys.stderr, 'Writing to stderr!'

if __name__ == '__main__':
  foo()
  colorFoo()
  foo()

This can still be polished a bit, but it should do the job in most cases and cleans up after itself properly. Of course, keep in mind I'm using shell-specific escape codes. If you want portability, you'll have to replace the escape code writes with calls to a portable terminal control module.

Answer 2

Here is my code with termcolor module:

from termcolor import colored

class ColoredOutput:
    def __init__(self, org_handle, color, on_color=None, attrs=['bold']):
        self.org_handle = org_handle
        def wrapper_write(x):
            return org_handle.write(colored(x, color=color, on_color=on_color, attrs=attrs))
        self.wrapper_write = wrapper_write
    def __getattr__(self, attr):
        return self.wrapper_write if attr == 'write' else getattr(self.org_handle, attr)

if __name__ == '__main__':
    import sys
    import colorama   # I'm working under windows 7, so i need this module to enable terminal color

    colorama.init()
    sys.stderr = ColoredOutput(sys.stderr, 'red')
    print('This is a test string', file=sys.stderr)