"從腳本中捕獲標准輸出？"

Question

假設有一個腳本在做這樣的事情：

# module writer.py
import sys

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

Answer 1

對於未來的訪問者：Python 3.4 上下文庫通過redirect_stdout上下文管理器直接提供（參見Python上下文庫幫助）：

from contextlib import redirect_stdout
import io

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

Answer 2

設置stdout是一種合理的方法。 另一種是將其作為另一個進程運行：

import subprocess

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

Answer 3

這是您的代碼的上下文管理器版本。 它產生一個包含兩個值的列表； 第一個是標准輸出，第二個是標准錯誤。

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 4

從 Python 3 開始，您還可以使用sys.stdout.buffer.write()將（已經）編碼的字節字符串寫入 stdout（請參閱Python 3 中的 stdout ）。 當你這樣做時，簡單的StringIO方法不起作用，因為sys.stdout.encoding和sys.stdout.buffer都不可用。

從 Python 2.6 開始，您可以使用TextIOBase API ，其中包括缺少的屬性：

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 5

這是我的原始代碼的裝飾器對應物。

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

Answer 6

或者也許使用已經存在的功能......

from IPython.utils.capture import capture_output

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

c()

print c.stdout

Answer 7

我認為你應該看看這四個對象：

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

Answer 8

我喜歡 contextmanager 解決方案，但是如果您需要使用打開的文件和 fileno 支持存儲的緩沖區，您可以執行類似的操作。

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

Answer 9

這里的問題（如何重定向輸出的示例，而不是tee部分）使用os.dup2在操作系統級別重定向流。 這很好，因為它也適用於您從程序中生成的命令。

Answer 10

這是一個上下文管理器，其靈感來自@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 11

from io import StringIO 
import threading
import sys
import os
import time

class Cap():
    def __init__(self):
        self._stdout = None
        self._stderr = None
        self._r = None
        self._w = None
        self._thread = None
        self._on_readline_cb = None
        
    def __enter__(self):
        self._stdout = sys.stdout
        self._stderr = sys.stderr
        r, w = os.pipe()
        r, w = os.fdopen(r, 'r'), os.fdopen(w, 'w', 1)
        self._r = r
        self._w = w
        sys.stdout = self._w
        sys.stderr = self._w
        self._thread = threading.Thread(target=self._handler)
        self._thread.start()
        return self
    
    def _handler(self):
        while not self._w.closed:
            try:
                while True:
                    line = self._r.readline()
                    if len(line) == 0: break
                    if self._on_readline_cb: self._on_readline_cb(line)
            except:
                break
    
    def print(self, s, end=""):
        print(s, file=self._stdout, end=end)
      
    def on_readline(self, callback):
        self._on_readline_cb = callback
        
    def __exit__(self, *args):
        self._w.close()
        if self._thread: self._thread.join()
        self._r.close()
        sys.stdout = self._stdout
        sys.stderr = self._stderr
        
with Cap() as output:
    def callback(line):
        output.print("callback=",line)
    output.on_readline(callback)
    print('hello world1')
    print('hello world2')

Answer 12

當第三方代碼已經復制了對sys.stdout<\/code>的引用時，另一種方法是臨時替換write()<\/code>方法本身：

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