(python) colour printing with decorator in a function
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开发者_StackOverflow specify the colours, defaulting them to red and green?
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.
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)
This works for me in Bash in my Mac's Terminal.app
import sys
green = '\033[01;32m'
red = '\033[01;31m'
sys.stdout.write(green+"Hello ")
sys.stderr.write(red+"world!")
精彩评论