Hier j’ai eu rencontré le travail d’une de ces fameuses personnes qui pensent que la ré-utilisabilité c’est pour les pédés, et qui font des scripts dont la moitié des infos renvoyées sont printées au milieu de blocs de code de 50 lignes, sans possibilité de les récupérer.

Heureusement, avec un petit hack, on peut capturer ce qu’affiche un autre code, et sauver le bébé, l’eau du bain, et même le canard en plastique.

Le code pour les gens pressés

J’ai enrobé l’astuce dans un context manager, ça rend l’utilisation plus simple.

import sys from io import BytesIO from contextlib import contextmanager @ contextmanager def capture_ouput ( stdout_to = None , stderr_to = None ) : try : stdout , stderr = sys . stdout , sys . stderr sys . stdout = c1 = stdout_to or BytesIO ( ) sys . stderr = c2 = stderr_to or BytesIO ( ) yield c1 , c2 finally : sys . stdout = stdout sys . stderr = stderr try : c1. flush ( ) c1. seek ( 0 ) except ( ValueError , IOError ) : pass try : c2. flush ( ) c2. seek ( 0 ) except ( ValueError , IOError ) : pass import sys from io import BytesIO from contextlib import contextmanager @contextmanager def capture_ouput(stdout_to=None, stderr_to=None): try: stdout, stderr = sys.stdout, sys.stderr sys.stdout = c1 = stdout_to or BytesIO() sys.stderr = c2 = stderr_to or BytesIO() yield c1, c2 finally: sys.stdout = stdout sys.stderr = stderr try: c1.flush() c1.seek(0) except (ValueError, IOError): pass try: c2.flush() c2.seek(0) except (ValueError, IOError): pass

Notez l’usage de yield.

Et ça s’utilise comme ça:

with capture_output ( ) as stdout , stderr: fonction_qui_fait_que_printer_la_biatch ( ) print stdout. read ( ) # on récupère le contenu des prints with capture_output() as stdout, stderr: fonction_qui_fait_que_printer_la_biatch() print stdout.read() # on récupère le contenu des prints

Attention, le code n’est pas thread safe, c’est fait pour hacker un code crade, pas pour devenir une institution. Mais c’est fort pratique dans notre cas précis.

Comment ça marche ?

stdin (entrée standard), stdout (sortie standard) et stderr (sortie des erreurs) sont des file like objects, c’est à dire qu’ils implémentent l’interface d’un objet fichier: on peut les ouvrir, les lire, y écrire et les fermer avec des méthodes portant le même nom et acceptant les mêmes paramètres.

L’avantage d’avoir une interface commune, c’est qu’on peut du coup échanger un file like objet par un autre.

Par exemple on peut faire ceci:

import sys log = open ( '/tmp/log' , 'w' ) sys . stdout = log # hop, on hijack la sortie standard print "Hello" log. close ( ) import sys log = open('/tmp/log', 'w') sys.stdout = log # hop, on hijack la sortie standard print "Hello" log.close()

Comme print écrit dans stdout , en remplaçant stdout par un fichier, print va du coup écrire dans le fichier.

Mais ce code est fort dangereux, car il remplace stdout de manière définitive. Du coup, si du code print après, il va écrire dans le fichier, même les libs externes, car stdout est le même pour tout le monde dans le process Python courant.

Du coup, il est de bon ton de s’assurer la restauration de stdout à son état d’origine:

import sys log = open ( '/tmp/log' , 'w' ) bak = sys . stdout # on sauvegarde l'ancien stdout sys . stdout = log print "Hello" log. close ( ) sys . stdout = bak # on restore stdout import sys log = open('/tmp/log', 'w') bak = sys.stdout # on sauvegarde l'ancien stdout sys.stdout = log print "Hello" log.close() sys.stdout = bak # on restore stdout

Comme je le disais plus haut, ceci n’est évidement pas thread safe, puisqu’entre la hijacking et la restoration de stdout , un autre thread peut faire un print .

Dans notre context manager, on utilise BytesIO() et non un fichier. BytesIO est un file like objet qui permet de récupérer un flux de bits en mémoire. Donc on fait écrire print dedans, ainsi on a tout ce qu’on affiche qui se sauvegarde en mémoire.

Bien entendu, vous pouvez créé vos propres file like objects, par exemple un objet qui affiche à l’écran ET capture la sortie. Par exemple, pour mitiger le problème de l’absence de thread safe: 99% des libs n’ont pas besoin du vrai stdout , juste d’un truc qui print .

import sys from io import BytesIO class PersistentStdout ( object ) : old_stdout = sys . stdout def __init__ ( self ) : self . memory = BytesIO ( ) def write ( self , s ) : self . memory . write ( s ) self . old_stdout . write ( s ) old_stdout = sys . stdout sys . stdout = PersistentStdout ( ) print "test" # ceci est capturé et affiché sys . stdout . memory . seek ( 0 ) res = sys . stdout . memory . read ( ) sys . stdout = PersistentStdout. old_stdout print res # résultat de la capture import sys from io import BytesIO class PersistentStdout(object): old_stdout = sys.stdout def __init__(self): self.memory = BytesIO() def write(self, s): self.memory.write(s) self.old_stdout.write(s) old_stdout = sys.stdout sys.stdout = PersistentStdout() print "test" # ceci est capturé et affiché sys.stdout.memory.seek(0) res = sys.stdout.memory.read() sys.stdout = PersistentStdout.old_stdout print res # résultat de la capture