Capturar stdout de um script?

94

suponha que haja um script fazendo algo assim:

# module writer.py
import sys

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

Agora, suponha que eu queira capturar a saída do write função e armazená-la em uma variável para processamento posterior. A solução ingênua foi:

# module mymodule.py
from writer import write

out = write()
print out.upper()

Mas isso não funciona. Eu encontrei outra solução e funciona, mas, por favor, me diga se existe uma maneira melhor de resolver o problema. obrigado

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
Paolo
fonte

Respostas:

51

A configuração stdouté uma maneira razoável de fazer isso. Outra é executá-lo como outro processo:

import subprocess

proc = subprocess.Popen(["python", "-c", "import writer; writer.write()"], stdout=subprocess.PIPE)
out = proc.communicate()[0]
print out.upper()
Matthew Flaschen
fonte
4
check_output captura diretamente a saída de um comando executado em um subprocesso: <br> value = subprocess.check_output (command, shell = True)
Arthur
1
Versão formatada :value = subprocess.check_output(command, shell=True)
Nae
51

Aqui está uma versão do gerenciador de contexto do seu código. Ele produz uma lista de dois valores; o primeiro é stdout, o segundo é 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'
Jason Grout
fonte
Adoro esta solução. Eu modifiquei, para não perder acidentalmente coisas de um fluxo no qual não estou esperando saída, por exemplo, erros inesperados. No meu caso, capture () pode aceitar sys.stderr ou sys.stdout como parâmetro, indicando para capturar apenas aquele fluxo.
Joshua Richardson
StringIO não oferece suporte a Unicode de forma alguma, então você pode integrar a resposta aqui para tornar o suporte acima de caracteres não ASCII: stackoverflow.com/a/1819009/425050
mafrosis
3
Modificar um valor gerado no final é realmente muito estranho - with capture() as out:vai se comportar de forma diferente dewith capture() as out, err:
Eric
1
O suporte Unicode / stdout.buffer pode ser alcançado usando o módulo io. Veja minha resposta .
JonnyJD
3
Esta solução quebra se você usar subprocesse redirecionar a saída para sys.stdout / stderr. Isso ocorre porque StringIOnão é um objeto de arquivo real e perde a fileno()função.
letmaik
47

Para visitantes futuros: o contextlib do Python 3.4 fornece isso diretamente (consulte a ajuda do contextlib do Python ) por meio do redirect_stdoutgerenciador de contexto:

from contextlib import redirect_stdout
import io

f = io.StringIO()
with redirect_stdout(f):
    help(pow)
s = f.getvalue()
nodesr
fonte
2
Isso não resolve o problema ao tentar gravar em sys.stdout.buffer (como você precisa fazer ao gravar bytes). StringIO não tem o atributo buffer, enquanto TextIOWrapper tem. Veja a resposta de @JonnyJD.
weaver de
10

Esta é a contraparte decoradora do meu código original.

writer.py continua o mesmo:

import sys

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

mymodule.py ligeiramente modificado:

from writer import write as _write
from decorators import capture

@capture
def write():
    return _write()

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

E aqui está o decorador:

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
Paolo
fonte
10

Ou talvez use uma funcionalidade que já existe ...

from IPython.utils.capture import capture_output

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

c()

print c.stdout
dgrigonis
fonte
9

A partir do Python 3, você também pode usar sys.stdout.buffer.write()para escrever (já) strings de byte codificadas para stdout (consulte stdout em Python 3 ). Quando você faz isso, a StringIOabordagem simples não funciona porque nem sys.stdout.encodingnem sys.stdout.bufferestaria disponível.

A partir do Python 2.6, você pode usar a TextIOBaseAPI , que inclui os atributos ausentes:

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

Esta solução funciona para Python 2> = 2.6 e Python 3. Observe que nosso sys.stdout.write()só aceita strings Unicode e sys.stdout.buffer.write()só aceita strings de byte. Esse pode não ser o caso do código antigo, mas geralmente é o caso do código criado para ser executado no Python 2 e 3 sem alterações.

Se você precisa oferecer suporte a um código que envia strings de bytes para stdout diretamente, sem usar stdout.buffer, pode usar esta variação:

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)

Você não precisa definir a codificação do buffer de sys.stdout.encoding, mas isso ajuda ao usar este método para testar / comparar a saída do script.

JonnyJD
fonte
4

A questão aqui (o exemplo de como redirecionar a saída, não a teeparte) usa os.dup2para redirecionar um fluxo no nível do sistema operacional. Isso é bom porque também se aplica a comandos que você cria de seu programa.

Jeremiah Willcock
fonte
4

Acho que você deve olhar para estes quatro objetos:

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

Exemplo:

from writer import write

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

UPD: Como Eric disse em um comentário, não se deve usá-los diretamente, então copiei e colei.

# 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")
Oleksandr Fedorov
fonte
4

Eu gosto da solução contextmanager, no entanto, se você precisar do buffer armazenado com o arquivo aberto e suporte a fileno, poderá fazer algo assim.

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

usar

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
Nathan Buckner
fonte
2

Aqui está um gerenciador de contexto que se inspira na resposta de @JonnyJD apoiando a gravação de bytes em bufferatributos, mas também aproveitando os referênios dunder-io de sys para simplificação posterior.

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

O resultado é:

stdout: foo

stderr: bar
Ryne Everett
fonte