Capturer stdout à partir d'un script?

Supposons qu'un script fasse quelque chose comme ceci:

# module writer.py
import sys

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

Supposons maintenant que je veuille capturer la sortie de la writefonction et la stocker dans une variable pour un traitement ultérieur. La solution naïve était:

# module mymodule.py
from writer import write

out = write()
print out.upper()

Mais ça ne marche pas. Je propose une autre solution et cela fonctionne, mais s'il vous plaît, faites-moi savoir s'il existe une meilleure façon de résoudre le problème. Merci

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

Le réglage stdoutest une manière raisonnable de le faire. Une autre consiste à l'exécuter comme un autre processus:

import subprocess

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

Voici une version de gestionnaire de contexte de votre code. Il donne une liste de deux valeurs; le premier est stdout, le second est 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'

Pour les futurs visiteurs: Python 3.4 contextlib le fournit directement (voir l' aide contextlib de Python ) via le redirect_stdoutgestionnaire de contexte:

from contextlib import redirect_stdout
import io

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

C'est la contrepartie décoratrice de mon code original.

writer.py reste le même:

import sys

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

mymodule.py est légèrement modifié:

from writer import write as _write
from decorators import capture

@capture
def write():
    return _write()

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

Et voici le décorateur:

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

Ou peut-être utiliser des fonctionnalités déjà présentes ...

from IPython.utils.capture import capture_output

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

c()

print c.stdout

À partir de Python 3, vous pouvez également utiliser sys.stdout.buffer.write()pour écrire (déjà) des chaînes d'octets encodées dans stdout (voir stdout dans Python 3 ). Lorsque vous faites cela, l' StringIOapproche simple ne fonctionne pas car ni sys.stdout.encodingnisys.stdout.buffer ne serait disponible.

À partir de Python 2.6, vous pouvez utiliser l' TextIOBaseAPI , qui comprend les attributs manquants:

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

Cette solution fonctionne pour Python 2> = 2.6 et Python 3. Veuillez noter que notre sys.stdout.write()n'accepte que les chaînes Unicode et sys.stdout.buffer.write()n'accepte que les chaînes d'octets. Ce n'est peut-être pas le cas pour l'ancien code, mais c'est souvent le cas pour le code qui est construit pour s'exécuter sur Python 2 et 3 sans modifications.

Si vous avez besoin de prendre en charge le code qui envoie des chaînes d'octets à stdout directement sans utiliser stdout.buffer, vous pouvez utiliser cette variante:

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)

Vous n'avez pas à définir l'encodage du tampon sys.stdout.encoding, mais cela aide lorsque vous utilisez cette méthode pour tester / comparer la sortie de script.

La question ici (l'exemple de la façon de rediriger la sortie, pas la teepartie) utilise os.dup2pour rediriger un flux au niveau du système d'exploitation. C'est bien car cela s'appliquera également aux commandes que vous générez à partir de votre programme.

Je pense que vous devriez regarder ces quatre objets:

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

Exemple:

from writer import write

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

UPD: Comme Eric l'a dit dans un commentaire, il ne faut pas les utiliser directement, alors je l'ai copié et collé.

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

J'aime la solution contextmanager, mais si vous avez besoin du tampon stocké avec le fichier ouvert et le support de fileno, vous pouvez faire quelque chose comme ça.

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

utilisation

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

Voici un gestionnaire de contexte s'inspirant de la réponse de @ JonnyJD prenant en charge l'écriture d'octets vers des bufferattributs, mais tirant également parti des références dunder-io de sys pour une simplification supplémentaire.

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

La sortie est:

stdout: foo

stderr: bar