Python 封装函数以打印到变量

Question

如果我有一个包含很多打印语句的函数：

比如：

def funA():
  print "Hi"
  print "There"
  print "Friend"
  print "!"

我想做的事情是这样的：

def main():
  ##funA() does not print to screen here
  a = getPrint(funA()) ##where getPrint is some made up function/object
  print a ##prints what funA would normally print at this step

所以当调用funcA的时候，它不会进行任何打印，而是将结果输出到一个对象中。然后我再打印这个对象来获取结果。有没有办法做到这一点？我也不想修改原来的函数。

希望这样说清楚了。

Answer 1

最好的方法是使用上下文管理器

from contextlib import contextmanager
import StringIO
import sys

@contextmanager
def capture():
    old_stdout = sys.stdout
    sys.stdout = StringIO.StringIO()
    try:
        yield sys.stdout
    finally:
        sys.stdout = old_stdout

现在你可以运行任何打印的代码：

with capture() as c:
    funA()
    funB()
    print 'HELLO!'

然后稍后：

print c.getvalue()

Answer 2

from cStringIO import StringIO

def getPrint(func, *args, **kwds):
  old_stdout = sys.stdout
  sys.stdout = StringIO()
  try:
    func(*args, **kwds)
  except:
    raise
  else:
    return sys.stdout.getvalue()
  finally:
    sys.stdout = old_stdout

#...
a = getPrint(funA) # notice no (), it is called by getPrint
print a.rstrip("\n") # avoid extra trailing lines

当然可以！请把你想要翻译的内容发给我，我会帮你用简单易懂的语言解释清楚。

Answer 3

你几乎可以做到你想要的，只要你不介意有一点小小的语法差别：

import cStringIO
import sys

def getPrint(thefun, *a, **k):
  savstdout = sys.stdout
  sys.stdout = cStringIO.StringIO()
  try:
    thefun(*a, **k)
  finally:
    v = sys.stdout.getvalue()
    sys.stdout = savstdout
  return v

这个小差别就是你必须调用 getPrint(funA)，而不是 getPrint(funA()) —— 也就是说，你必须传递函数对象本身，而不是在后面加上小括号，这样就会立即调用它，在 getPrint 能发挥作用之前。

如果你真的坚持要加上那些额外的小括号，那么 getPrint 就无法完成所有必要的准备工作，还需要其他代码来正确地准备好（我强烈建议去掉那些多余的小括号，这样就能把所有功能都封装在 getPrint 里面！）。