Python 从脚本中捕获标准输出？

假设有一个脚本执行以下操作：

# 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

---

问题（如何重定向输出的示例，而不是

tee

部分）使用

os.dup2

在os级别重定向流。这很好，因为它也将应用于从程序生成的命令。

设置

stdout

是一种合理的方法。另一种方法是将其作为另一个进程运行：

import subprocess

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

---

这是与我的原始代码相对应的decorator

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

---

这是您的代码的上下文管理器版本。它产生一个包含两个值的列表；第一个是标准输出，第二个是标准输出

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（请参阅）。 这样做时，简单的

StringIO

方法不起作用，因为

sys.stdout.encoding

和

sys.stdout.buffer

都不可用

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

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())

此解决方案适用于Python2>=2.6和Python3。 请注意，我们的

sys.stdout.write（）

只接受unicode字符串，而

sys.stdout.buffer.write（）

只接受字节字符串。 对于旧代码来说可能不是这样，但是对于构建为在Python 2和3上运行而不做任何更改的代码来说通常是这样

如果您需要支持不使用stdout.buffer而直接将字节字符串发送到stdout的代码，则可以使用以下变体：

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.encoding，但这有助于使用此方法测试/比较脚本输出。

或者使用已有的功能

from IPython.utils.capture import capture_output

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

c()

print c.stdout

---

我喜欢contextmanager解决方案，但是如果您需要存储在打开的文件和文件中的缓冲区，那么您可以这样做

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()

使用

---

对于未来的访问者：Python 3.4 contextlib通过

重定向\u stdout

上下文管理器直接提供（请参阅）：

from contextlib import redirect_stdout
import io

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

---

这里有一个上下文管理器，它的灵感来自@JonnyJD支持将字节写入

buffer

属性，并利用它进一步简化

导入io
导入系统
导入上下文库
@contextlib.contextmanager
def捕获_输出（）：
输出={}
尝试：
#重定向
sys.stdout=io.TextIOWrapper（io.BytesIO（），sys.stdout.encoding）
sys.stderr=io.TextIOWrapper（io.BytesIO（），sys.stderr.encoding）
产量
最后：
#阅读
系统标准寻道（0）
系统标准寻道（0）
输出['stdout']=sys.stdout.read（）
输出['stderr']=sys.stderr.read（）
sys.stdout.close（）
sys.stderr.close（）
#恢复
sys.stdout=sys.\u stdout__
sys.stderr=sys.\u stderr__
将capture_output（）作为输出：
打印（'foo'）
sys.stderr.buffer.write（b'bar'）
打印（'stdout:{stdout}'。格式（stdout=output['stdout']））
打印（'stderr:{stderr}'。格式（stderr=output['stderr']））

输出为：

stdout: foo

stderr: bar

---

我喜欢这个解决方案。我进行了修改，以避免意外丢失我不希望输出的流中的内容，例如意外错误。在我的例子中，capture（）可以接受sys.stderr或sys.stdout作为参数，表示仅捕获该stream.StringIO不以任何方式支持unicode，因此您可以在这里集成答案以使上述支持非ASCII字符：使用capture（）修改finally中生成的值实际上是wierd-

as out:

的行为将不同于捕获（）为out的

，错误：

Unicode/stdout.buffer支持可以通过使用io模块来实现。请参阅。如果使用

子流程

并将输出重定向到sys.stdout/stderr，则此解决方案将中断。这是因为

StringIO

不是真正的文件对象，并且错过了

fileno（）

函数。check\u output直接捕获子流程中运行的命令的输出：
value=subprocess。check\u output（command，shell=True）在现代Python版本上，您可以执行

capture\u output=True

而不是

stdout=subprocess.PIPE

。这并不能解决尝试写入sys.stdout.buffer时的问题（就像写入字节时需要做的那样）。StringIO不具有缓冲区属性，而TextIOWrapper具有缓冲区属性。请参阅@JonnyJD的答案。

stdout: foo

stderr: bar