从脚本捕获 stdout？

假设有这样一个脚本:

# module writer.py
import sys


def write():
sys.stdout.write("foobar")

现在假设我想捕获 write函数的输出并将其存储在一个变量中以便进一步处理。天真的解决办法是:

# module mymodule.py
from writer import write


out = write()
print out.upper()

但这不管用。我想出了另一个解决方案，它的工作，但请让我知道，如果有一个更好的方法来解决这个问题。谢谢

import sys
from cStringIO import StringIO


# setup the environment
backup = sys.stdout


# ####
sys.stdout = StringIO()     # capture output
write()
out = sys.stdout.getvalue() # release output
# ####


sys.stdout.close()  # close the stream
sys.stdout = backup # restore original stdout


print out.upper()   # post processing
119366 次浏览
小开

问题 给你(如何重定向输出的示例，而不是 tee部分)使用 os.dup2在操作系统级别重定向流。这很好，因为它也将应用于从程序中产生的命令。

小开
最佳答案

设置 stdout是一种合理的方法，另一种方法是将它作为另一个进程运行:

import subprocess


proc = subprocess.Popen(["python", "-c", "import writer; writer.write()"], stdout=subprocess.PIPE)
out = proc.communicate()[0]
print out.upper()
小开

这是我原始代码的装饰器对应物。

writer.py保持不变:

import sys


def write():
sys.stdout.write("foobar")

mymodule.py稍作修改:

from writer import write as _write
from decorators import capture


@capture
def write():
return _write()


out = write()
# out post processing...

这位是室内设计师:

def capture(f):
"""
Decorator to capture standard output
"""
def captured(*args, **kwargs):
import sys
from cStringIO import StringIO


# setup the environment
backup = sys.stdout


try:
sys.stdout = StringIO()     # capture output
f(*args, **kwargs)
out = sys.stdout.getvalue() # release output
finally:
sys.stdout.close()  # close the stream
sys.stdout = backup # restore original stdout


return out # captured output wrapped in a string


return captured
小开

下面是代码的上下文管理器版本。它生成两个值的列表; 第一个是 stdout，第二个是 stderr。

import contextlib
@contextlib.contextmanager
def capture():
import sys
from cStringIO import StringIO
oldout,olderr = sys.stdout, sys.stderr
try:
out=[StringIO(), StringIO()]
sys.stdout,sys.stderr = out
yield out
finally:
sys.stdout,sys.stderr = oldout, olderr
out[0] = out[0].getvalue()
out[1] = out[1].getvalue()


with capture() as out:
print 'hi'
小开

我觉得你应该看看这四件东西:

from test.test_support import captured_stdout, captured_output, \
captured_stderr, captured_stdin

例如:

from writer import write


with captured_stdout() as stdout:
write()
print stdout.getvalue().upper()

UPD: 正如 Eric 在评论中所说，一个人不应该直接使用它们，所以我复制并粘贴了它。

# Code from test.test_support:
import contextlib
import sys


@contextlib.contextmanager
def captured_output(stream_name):
"""Return a context manager used by captured_stdout and captured_stdin
that temporarily replaces the sys stream *stream_name* with a StringIO."""
import StringIO
orig_stdout = getattr(sys, stream_name)
setattr(sys, stream_name, StringIO.StringIO())
try:
yield getattr(sys, stream_name)
finally:
setattr(sys, stream_name, orig_stdout)


def captured_stdout():
"""Capture the output of sys.stdout:


with captured_stdout() as s:
print "hello"
self.assertEqual(s.getvalue(), "hello")
"""
return captured_output("stdout")


def captured_stderr():
return captured_output("stderr")


def captured_stdin():
return captured_output("stdin")
小开

从 Python3开始，您还可以使用 sys.stdout.buffer.write()将(已经)编码的字节字符串写入 stdout (参见 Python 3中的 stdout)。 当您这样做时，简单的 StringIO方法不起作用，因为 sys.stdout.encoding和 sys.stdout.buffer都不可用。

从 Python 2.6开始，您可以使用 TextIOBaseAPI，其中包括缺少的属性:

import sys
from io import TextIOWrapper, BytesIO


# setup the environment
old_stdout = sys.stdout
sys.stdout = TextIOWrapper(BytesIO(), sys.stdout.encoding)


# do some writing (indirectly)
write("blub")


# get output
sys.stdout.seek(0)      # jump to the start
out = sys.stdout.read() # read output


# restore stdout
sys.stdout.close()
sys.stdout = old_stdout


# do stuff with the output
print(out.upper())

This solution works for Python 2 >= 2.6 and Python 3. 请注意，我们的 sys.stdout.write()只接受 unicode 字符串，而 sys.stdout.buffer.write()只接受字节字符串。 This might not be the case for old code, but is often the case for code that is built to run on Python 2 and 3 without changes.

If you need to support code that sends byte strings to stdout directly without using stdout.buffer, you can use this variation:

class StdoutBuffer(TextIOWrapper):
def write(self, string):
try:
return super(StdoutBuffer, self).write(string)
except TypeError:
# redirect encoded byte strings directly to buffer
return super(StdoutBuffer, self).buffer.write(string)

您不必将缓冲区的编码设置为 sys.stdout.coding，但是在使用这种方法测试/比较脚本输出时，这会有所帮助。

小开

或者使用已经存在的功能..。

from IPython.utils.capture import capture_output


with capture_output() as c:
print('some output')


c()


print c.stdout
小开

我喜欢上下文管理器的解决方案，但是，如果你需要的缓冲区存储与打开文件和文件支持，你可以这样做。

import six
from six.moves import StringIO




class FileWriteStore(object):
def __init__(self, file_):
self.__file__ = file_
self.__buff__ = StringIO()


def __getattribute__(self, name):
if name in {
"write", "writelines", "get_file_value", "__file__",
"__buff__"}:
return super(FileWriteStore, self).__getattribute__(name)
return self.__file__.__getattribute__(name)


def write(self, text):
if isinstance(text, six.string_types):
try:
self.__buff__.write(text)
except:
pass
self.__file__.write(text)


def writelines(self, lines):
try:
self.__buff__.writelines(lines)
except:
pass
self.__file__.writelines(lines)


def get_file_value(self):
return self.__buff__.getvalue()

使用

import sys
sys.stdout = FileWriteStore(sys.stdout)
print "test"
buffer = sys.stdout.get_file_value()
# you don't want to print the buffer while still storing
# else it will double in size every print
sys.stdout = sys.stdout.__file__
print buffer
小开

对于未来的访问者: Python 3.4 contextlib 通过 redirect_stdout上下文管理器直接提供了这个功能(参见 Python contextlib 帮助) :

from contextlib import redirect_stdout
import io


f = io.StringIO()
with redirect_stdout(f):
help(pow)
s = f.getvalue()
小开

下面是一个上下文管理器，它从@JonnyJD 支持将字节写入 buffer属性的 回答中获得灵感，同时也利用了 Sys 的 Dunder-io Referens来进一步简化。

import io
import sys
import contextlib




@contextlib.contextmanager
def capture_output():
output = {}
try:
# Redirect
sys.stdout = io.TextIOWrapper(io.BytesIO(), sys.stdout.encoding)
sys.stderr = io.TextIOWrapper(io.BytesIO(), sys.stderr.encoding)
yield output
finally:
# Read
sys.stdout.seek(0)
sys.stderr.seek(0)
output['stdout'] = sys.stdout.read()
output['stderr'] = sys.stderr.read()
sys.stdout.close()
sys.stderr.close()


# Restore
sys.stdout = sys.__stdout__
sys.stderr = sys.__stderr__




with capture_output() as output:
print('foo')
sys.stderr.buffer.write(b'bar')


print('stdout: {stdout}'.format(stdout=output['stdout']))
print('stderr: {stderr}'.format(stderr=output['stderr']))

产出为:

stdout: foo


stderr: bar
小开

当第三方代码已经复制了对 sys.stdout的引用时，另一种方法是临时替换 write()方法本身:

from types import MethodType
...
f = io.StringIO()
def new_write(self, data):
f.write(data)


old_write = sys.stdout.write
sys.stdout.write = MethodType(new_write, sys.stdout)
error = command.run(args)
sys.stdout.write = old_write
output = f.getvalue()

提交答案