Захватить стандартный вывод из скрипта на Python - PullRequest
75 голосов
/ 28 февраля 2011

предположим, что есть скрипт, выполняющий что-то вроде этого:

# 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

Ответы [ 9 ]

43 голосов
/ 28 февраля 2011

Установка stdout - разумный способ сделать это. Другой - запустить его как другой процесс:

import subprocess

proc = subprocess.Popen(["python", "-c", "import writer; writer.write()"], stdout=subprocess.PIPE)
out = proc.communicate()[0]
print out.upper()
41 голосов
/ 24 мая 2012

Вот версия контекстного менеджера вашего кода.Это дает список из двух значений;первый - стандартный вывод, второй - стандартный вывод.

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'
26 голосов
/ 04 ноября 2016

Для будущих посетителей: Python 3.4 contextlib предоставляет это напрямую (см. Справка Python contextlib ) через менеджер контекста redirect_stdout:

f = io.StringIO()
with redirect_stdout(f):
    help(pow)
s = f.getvalue()
10 голосов
/ 25 февраля 2016

Или, может быть, использовать функциональность, которая уже есть ...

from IPython.utils.capture import capture_output

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

c()

print c.stdout
9 голосов
/ 16 марта 2011

Это копия моего оригинального кода.

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
6 голосов
/ 13 октября 2013

Начиная с Python 3, вы также можете использовать sys.stdout.buffer.write() для записи (уже) закодированных байтовых строк в стандартный вывод (см. стандартный вывод в Python 3 ). Когда вы делаете это, простой 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.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 голосов
/ 11 октября 2016

Мне нравится решение 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()

use

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
3 голосов
/ 17 сентября 2013

Думаю, вам стоит взглянуть на эти четыре объекта:

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: Как сказал Эрик в комментарии, их не следует использовать напрямую, поэтому я скопировал и вставил их.

# 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 голосов
/ 28 февраля 2011

Вопрос здесь (пример того, как перенаправить вывод, а не часть tee) использует os.dup2 для перенаправления потока на уровне ОС. Это хорошо, потому что это будет применяться и к командам, которые вы порождаете из вашей программы.

...