Capturer l’affichage des prints d’un code Python
samedi 29 septembre 2012 à 16:03Hier 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
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
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()
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
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
Pour cette raison le code du context manager permet de passer le file like objet à utiliser en argument. On notera aussi que si on souhaite rediriger stdout
mais pas stderr
et vice-versa, il suffit de passer sys.stdout
et sys.stderr
en argument :-)