（python）在函数中使用装饰器进行彩色打印

Question

我想给一个函数加个装饰，让它打印到 stdout 的内容是绿色，而打印到 stderr 的内容是红色。我手头有 termcolor 这个模块可以用。

额外加分：我该怎么做才能给这个装饰器传递参数，以指定颜色，默认颜色是红色和绿色？

Answer 1

在我Mac的Terminal.app中，这在Bash里对我有效。

import sys
green = '\033[01;32m'
red = '\033[01;31m'

sys.stdout.write(green+"Hello ")
sys.stderr.write(red+"world!")

Answer 2

这是我使用termcolor模块的代码：

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)

Answer 3

这是个有趣的问题。最简单的解决办法就像Pete提到的那样。在运行函数之前，先打印出一些特殊的代码到标准错误输出和标准输出。不过，如果这两个输出都指向同一个终端（通常情况下就是这样），它们会互相干扰。

所以，另一种解决方案是对标准输出和标准错误输出进行“猴子补丁”，也就是用一个小的包装器来在每次写入时启用颜色，这样做的时候要确保我们是在终端中（而不是通过管道）。

#!/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()

这个方法还可以进一步优化，但在大多数情况下应该能正常工作，并且会妥善清理自己。当然，要记住我使用的是特定于shell的特殊代码。如果你想要更好的兼容性，就需要把这些特殊代码替换成一个通用的终端控制模块的调用。