从脚本捕获标准输出?


89

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

# 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

Answers:


49

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

import subprocess

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

4
check_output直接捕获在子进程中运行的命令的输出:<br> value = subprocess.check_output(command,shell = True)
Arthur

1
格式化的版本value = subprocess.check_output(command, shell=True)
Nae

45

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

喜欢这个解决方案。我进行了修改,以免意外丢失我不期望输出的流中的内容,例如意外错误。在我的情况下,capture()可以接受sys.stderr或sys.stdout作为参数,指示仅捕获该流。
约书亚·理查森

StringIO不以任何方式支持unicode,因此您可以在此处集成答案以使上面的内容支持非ASCII字符: stackoverflow.com/a/1819009/425050
mafrosis

2
最终修改一个产生的值确实很奇怪-with capture() as out:行为与with capture() as out, err:
Eric

可以使用io模块获得Unicode / stdout.buffer支持。看我的回答
JonnyJD

1
如果您使用subprocess输出并将其重定向到sys.stdout / stderr,则此解决方案会中断。这是因为StringIO它不是真实的文件对象,并且缺少该fileno()功能。
letmaik 2014年

44

对于未来的访问者:Python 3.4 contextlib通过上下文管理器直接提供了此功能(请参阅Python contextlib helpredirect_stdout

from contextlib import redirect_stdout
import io

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

尝试写入sys.stdout.buffer时,这不能解决问题(如写入字节时需要这样做)。StringIO没有buffer属性,而TextIOWrapper有。请参阅@JonnyJD的答案。
韦弗

9

这是我原始代码的装饰器副本。

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

9

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

from IPython.utils.capture import capture_output

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

c()

print c.stdout

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和Python3。请注意,我们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,但这在使用此方法测试/比较脚本输出时会有所帮助。



3

我认为您应该查看以下四个对象:

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

3

我喜欢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()

使用

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

0

这是一个上下文管理器,它从@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
By using our site, you acknowledge that you have read and understand our Cookie Policy and Privacy Policy.
Licensed under cc by-sa 3.0 with attribution required.