Chụp stdout từ một tập lệnh?

giả sử có một tập lệnh làm điều gì đó như thế này:

# module writer.py
import sys

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

Bây giờ, giả sử tôi muốn nắm bắt đầu ra của writehàm và lưu trữ nó trong một biến để xử lý thêm. Giải pháp ngây thơ là:

# module mymodule.py
from writer import write

out = write()
print out.upper()

Nhưng điều này không hiệu quả. Tôi đưa ra một giải pháp khác và nó hoạt động, nhưng xin vui lòng cho tôi biết nếu có cách nào tốt hơn để giải quyết vấn đề. Cảm ơn

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

Đặt stdoutlà một cách hợp lý để làm điều đó. Cách khác là chạy nó như một quá trình khác:

import subprocess

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

Đây là phiên bản trình quản lý ngữ cảnh của mã của bạn. Nó tạo ra một danh sách hai giá trị; đầu tiên là stdout, thứ hai là 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'

Đối với khách truy cập trong tương lai: Trình quản lý ngữ cảnh Python 3.4 cung cấp trực tiếp cho điều này (xem trợ giúp về redirect_stdoutNgữ cảnh Python ) thông qua trình quản lý ngữ cảnh:

from contextlib import redirect_stdout
import io

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

Đây là bản sao trang trí của mã gốc của tôi.

writer.py vẫn như cũ:

import sys

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

mymodule.py sligthly được sửa đổi:

from writer import write as _write
from decorators import capture

@capture
def write():
    return _write()

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

Và đây là người trang trí:

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

Hoặc có thể sử dụng chức năng đã có ...

from IPython.utils.capture import capture_output

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

c()

print c.stdout

Bắt đầu với Python 3, bạn cũng có thể sử dụng sys.stdout.buffer.write()để viết các chuỗi byte được mã hóa (đã có) vào stdout (xem stdout trong Python 3 ). Khi bạn làm điều đó, StringIOcách tiếp cận đơn giản không hoạt động vì cả hai đều sys.stdout.encodingkhông sys.stdout.buffercó sẵn.

Bắt đầu với Python 2.6, bạn có thể sử dụng TextIOBaseAPI , bao gồm các thuộc tính bị thiếu:

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

Giải pháp này hoạt động cho Python 2> = 2.6 và Python 3. Xin lưu ý rằng chúng tôi sys.stdout.write()chỉ chấp nhận chuỗi unicode và sys.stdout.buffer.write()chỉ chấp nhận chuỗi byte. Đây có thể không phải là trường hợp đối với mã cũ, nhưng thường là trường hợp đối với mã được xây dựng để chạy trên Python 2 và 3 mà không có thay đổi.

Nếu bạn cần hỗ trợ mã gửi chuỗi byte tới stdout trực tiếp mà không sử dụng stdout.buffer, bạn có thể sử dụng biến thể này:

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)

Bạn không phải đặt mã hóa của bộ đệm là sys.stdout.encoding, nhưng điều này sẽ hữu ích khi sử dụng phương pháp này để kiểm tra / so sánh đầu ra tập lệnh.

Câu hỏi ở đây (ví dụ về cách chuyển hướng đầu ra, không phải teephần) sử dụng os.dup2để chuyển hướng một luồng ở cấp hệ điều hành. Điều đó rất hay vì nó cũng sẽ áp dụng cho các lệnh mà bạn sinh ra từ chương trình của mình.

Tôi nghĩ Bạn nên xem xét bốn đối tượng sau:

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

Thí dụ:

from writer import write

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

UPD: Như Eric đã nói trong một bình luận, người ta không nên sử dụng chúng trực tiếp, vì vậy tôi đã sao chép và dán nó.

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

Tôi thích giải pháp trình quản lý ngữ cảnh, tuy nhiên nếu bạn cần bộ đệm được lưu trữ cùng với tệp đang mở và hỗ trợ tệpno thì bạn có thể làm như thế này.

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

sử dụng

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

Dưới đây là một trình quản lý ngữ cảnh lấy cảm hứng từ câu trả lời của @ JonnyJD hỗ trợ việc viết byte vào bufferthuộc tính abut cũng tận dụng các tham chiếu dunder-io của sys để đơn giản hóa hơn nữa.

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']))

Đầu ra là:

stdout: foo

stderr: bar