Supposons que vous vouliez afficher quelque chose avant un bout de code, et après un bout de code, même si celui-ci rate. Vous feriez quelque chose comme ça:

def truc():
    print "machin"
 
print "Avant"
try:
    truc()
finally:
    print "Après"

Et ça va afficher:

Avant
machin
Après

Et avec:

def truc():
    print "machin"
    raise Exception('Fail !')

‘Après’ sera quand même affiché. Ça plantera, mais la dernière action sera toujours faite.

Si vous le faites souvent, vous voudrez factoriser du code. Un des moyens de le faire est d’utiliser les context managers.

Créer son propre context manager

Un context manager est une classe ordinaire en Python. Sa seule spécificité est de déclarer une méthode __enter__() et une méthode __exit__(). Ces méthodes sont des méthodes ordinaires, leur nom spécial est juste là par convention, et en les nommant ainsi on s’assure qu’elles seront détectées et utilisées automatiquement.

Notre code là haut peut donc se réécrire ainsi:

class MonSuperContextManager(object):
    def __enter__(self):
        print "Avant"
    def __exit__(self, type, value, traceback):
        # faites pas attention aux paramètres, ce sont toutes les infos
        # automatiquement passées à __exit__ et qui servent pour inspecter
        # une éventuelle exception
        print "Après"
 
with MonSuperContextManager():
    truc()

L’avantage de with est multiple:

• Il permet de visualiser très précisément où on entre dans l’action et où on en sort (c’est un seul block)
• Il permet de réutiliser les actions faite à l’entrée et à la sortie de l’action.
• Même si une exception est levée, l’action de sortie sera exécutée juste avant le plantage. __exit__ est en effet garantie d’être appelée quoiqu’il arrive. Bon, évidement, si il y a une coupure de courant…

En gros, créer un context manager, c’est faire un raccourci lisible pour try/finally. Point.

Un exemple utile de context manager

Supposons que vous ayez beaucoup de travail à faire dans plein de dossiers. Vous voulez vous assurer que vous allez dans le dossier de travail, puis que vous retournez au dossier initial à chaque fois.

import os
 
class Cd(objet):
    def __init__(dirname):
        self.dirname = dirname
    def __enter__(self):
        self.curdir = os.getcwd()
        os.chdir(self.dirname)
    def __exit__(self, type, value, traceback):
        os.chdir(self.curdir)

On l’utilise comme ça:

# ici on est dans /home/moi
 
with Cd('/'):
 
    # faire un truc dans /
 
    with Cd('/opt'):
 
        # faire un truc dans /opt
 
    # ici on est dans /
 
# ici on est dans /home/moi

C’est d’ailleurs ce que fait fabric.

Le mot clé as

Tout ce qu’on retourne dans __enter__ peut être récupéré grâce au mot clé as. Imaginons un context manager qui permette d’ouvrir un fichier et de le fermer automatiquement:

class OpenFile(objet):
    def __init__(filename, mode='r'):
        self.filename = filename
        self.mode = mode
    def __enter__(self):
        self.file = open(self.filename, self.mode)
        # ici on retourne l'objet fichier, il sera accessible avec "as"
        return self.file
    def __exit__(self, type, value, traceback):
        self.file.close()

On l’utilise comme ceci:

with OpenFile('/etc/fstab') as f:
    for line in f:
        print line

f va contenir ici l’objet fichier, car nous l’avons retourné dans __enter__. A la fin du bloc with, le fichier sera fermé automatiquement.

Et devinez quoi, Python possède déjà un context manager qui fait ça:.

with open(vot_fichier_msieu_dames) as f:
   # faire un truc

Context managers sous forme de fonctions

Faire les choses sous forme de classes, c’est pratique quand on a beaucoup de logique à encapsuler. Mais la plupart des context managers sont très simples. Pour cette raison, Python vient avec plein d’outils pour se simplifier la vie avec with dans un module judicieusement nommé contextlib.

Pour l’utiliser, il faut avoir des notions sur les décorateurs, et le mot clé yield. Si ce n’est pas votre cas, restez sur la version sous forme de classe :-)

Supposons que l’on veuille recréer le context manager open:

from contextlib import contextmanager
 
@contextmanager
def open(filename, mode):
    try:
        f = open(filename, mode)
        yield f
    finally:
        f.close()

Bon, c’est simplifié, hein, le vrai est plus robuste que ça.

Comment ça marche ?

D’abord, on utilise le décorateur @contextmanager pour dire à Python que la fonction sera un context manager.

Ensuite, on fait un try/finally (il est pas automatique comme avec __enter__ et __exit__).

yield sépare le code en deux: tout ce qui est avant est l’équivalent de __enter__, tout ce qui est après est l’équivalent de __exit__. Ce qui est “yieldé” est ce que l’on récupère avec le mot clé as.

Context manager et décorateur, le shampoing deux en un

Ces deux fonctionnalités se ressemblent beaucoup: elles permettent toutes les deux de lancer du code automatiquement avant et après un code tiers. La seule différence est que le context manager le fait à la demande, alors que le décorateur s’applique à la définition d’une fonction.

Quand on sait comment ils marchent, il est facile de faire un context manager utilisable également en tant que décorateur.

from functools import wraps
 
class ContextDecorator(object):
    # __call__ est une méthode magique appelée quand on utilise () sur un objet
    def __call__(self, f):
        # bon, cette partie là suppose que vous savez comment marche un
        # décorateur, si c'est pas le cas, retournez lire l'article sur S&M
        # linké dans le premier paragraphe
        @wraps(f)
        def decorated(*args, **kwds):
            # notez le with appelé sur soi-même, c'est y pas mignon !
            with self:
                return f(*args, **kwds)
        return decorated

Et voilà, il suffit d’hériter de ça, et on a un décorateur + context manager. Par exemple, si on veut timer un truc:

import datetime
 
class TimeIt(ContextDecorator):
 
    def __enter__(self):
        self.start = datetime.datetime.now()
        print self.start
 
    def __exit__(self, type, value, traceback):
        print (datetime.datetime.now() -self.start).total_seconds()

Timer juste un appel:

def foo():
    # faire un truc
 
with TimeIt():
    foo()

Timer tous les appels:

@TimeIt()
def foo():
   # faire un truc

Notez que ContextDecorator est présent par défaut dans le module contextlib sous Python 3.2.

Prenez un model Django Media: vous transcodez ce media, vous faites souvent des opérations dans son dossier. Il faut s’assurer que ça marche. Et quand ça foire, il faut s’assurer que l’on a bien nettoyé le répertoire de travail.

Hop:

class Media(models.Model):
 
    [...]
 
    @contextmanager
    def in_dir(self, delete_on_error=False):
        # si vous ne comprenez pas ce code, ne vous en faites pas,
        # c'est du Python avancé
        try:
            os.makedirs(self.dir)
        except (OSError, IOError) as e:
            pass
 
        try:
            yield self.dir
        except:
            if delete_on_error:
                shutil.rmtree(self.dir)
            raise

Et voilà, Max dans ses scripts de transcoding magiques n’a pas à se poser de question. Il a juste à faire:

with media.in_dir(delete_on_error=True) as d:
    # faire des trucs avec le dossier dont le chemin est dans "d"

Le chemin spécifique pour ce média est garanti d’exister.

Si il y a une exception, le dossier est garanti d’être nettoyé.

C’est simple, c’est propre, c’est beau. J’aime.