从Python中的脚本捕获stdout

Question

假设有一个脚本做这样的事情:

# 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

Answer 1

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

import subprocess

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

Answer 2

这是代码的上下文管理器版本.它产生两个值的列表; 第一个是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'

Answer 3

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

from contextlib import redirect_stdout
import io

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

Answer 4

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

from IPython.utils.capture import capture_output

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

c()

print c.stdout

Answer 5

这是我原始代码的装饰者对应物.

writer.py 保持原样:

import sys

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

mymodule.py sligthly被修改:

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

Answer 6

这是一个上下文管理器，它从 @JonnyJD 的答案中汲取灵感，支持将字节写入buffer属性，但也利用sys 的 dunder-io 引用来进一步简化。

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

Answer 7

从Python 3开始,您还可以使用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())

此解决方案适用于Python 2> = 2.6和Python 3.请注意,我们sys.stdout.write()只接受unicode字符串,sys.stdout.buffer.write()只接受字节字符串.对于旧代码可能不是这种情况,但对于构​​建为在Python 2和3上运行而无需更改的代码通常就是这种情况.

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

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,但这在使用此方法测试/比较脚本输出时会有所帮助.

Answer 8

这里的问题（如何重定向输出的示例，而不是部分tee）用于os.dup2在操作系统级别重定向流。这很好，因为它也适用于您从程序中生成的命令。