Context Manager

Context manager (kontextový manažer) pomáhá soustředit se na konkrétní úlohu nad blokem kódu a neřešit vedlejší, ale přesto důležité operace, jako je alokace nebo uvolnění zdrojů. Pokud chceme, například, něco zapsat do souboru, takový soubor musí být otevřen a po zápisu správně uvolněn, což je operace, na kterou může vývojář snadno zapomenout. Kontextový manažer soubor sám správně otevře a po opuštění with bloku uvolní. Dalším příkladem může být práce nad databází, kdy chceme v bloku kódu spustit SQL dotaz a kontextový manažer se nám postará o commit, rollback, případně uvolnění spojení.

Klíčové slovo with

Začneme příkladem, ve kterém budeme chtít provést zápis do souboru. Nejdříve bez použití kontextových manažerů:

f = open("soubor.txt", "w")
try:
    f.write("Ahoj světe")
finally:
    f.close()

Po zápisu bychom soubor měli uzavřít, uvolnit, aby do něj mohl zapisovat někdo jiný. Tj. kromě samotného zápisu musíme řešit ještě další režii. Zkusíme to teď pomocí kontextových manažerů a klíčového slova with:

with open("soubor.txt", "w") as f:
    f.write("Ahoj světe!")

Vidíme, že zápis se zjednodušil a zároveň máme jistotu, že v okamžiku, kdy program opustí with blok, dojde k uvolnění souboru. Jak to tedy funguje uvnitř?

Implementace pomocí třídy

Pokud potřebujeme napsat vlastní kontextový manažer, můžeme sáhnout po implementaci pomocí třídy. Taková třída musí implementovat určité rozhraní, konkrétně funkce __enter__() a __exit__().

class FileManager:

    def __init__(self, path, mode):
        print("__init__")
        self.path = path
        self.mode = mode
        self.file = None

    def __enter__(self):
        print("__enter__")
        self.file = open(self.path, self.mode)
        return self.file

    def __exit__(self, exc_type, exc_value, exc_traceback):
        print("__exit__")
        self.file.close()


print("Před blokem")

with FileManager("soubor.txt", "w") as f:
        print("Před zápisem do souboru")
        f.write("Ahoj světe!")
        print("Po zápisu do souboru")

print("Po bloku")

Výstup:

Před blokem
__init__
__enter__
Před zápisem do souboru
Po zápisu do souboru
__exit__
Po bloku

Za klíčovým slovem with následuje vytvoření instance třídy FileManager(). Po vytvoření instance a zavolání funkce __init__ se automaticky zavolá funkce __enter__. Tato funkce může vracet nějakou návratovou hodnotu, kterou si ve with bloku pomocí klíčového slova as můžeme načíst do nějaké lokální proměnné. Po opuštění with bloku se automaticky zavolá funkce __exit__. Můžeme říci, že kontext je v tomto případě definován vnitřním stavem proměnných třídy.

Výjimky

Co když ve with bloku dojde k nějaké výjimce? Co se stane s alokovanými zdroji? Uvolní je sám kontext manažer, nebo v tomto případě bude muset nějak zasáhnout vývojář přímo v tomto bloku?

class FileManager:

    def __init__(self, path, mode):
        self.path = path
        self.mode = mode
        self.file = None

    def __enter__(self):
        self.file = open(self.path, self.mode)
        return self.file

    def __exit__(self, exc_type, exc_value, exc_traceback):
        print(exc_type)
        self.file.close()

with FileManager("soubor.txt", "w") as f:
        f.write("Ahoj světe!")
        raise Exception("Něco se stalo!")

V kódu výše došlo k vyhození výjimky. Po spuštění programu dostaneme tento výstup:

<class 'Exception'>
Traceback (most recent call last):
  File "context-managers.py", line 18, in <module>
    raise Exception("Něco se stalo!")
Exception: Něco se stalo!

Kromě samotné výjimky je vidět, že ve funkci __exit__ je naplněna proměnná exc_type, ve které se nachází typ vyhozené výjimky. Platí tedy, že pokud exc_type není None, došlo k chybě a jsme schopni vykonat potřebnou operaci. Často je žádoucí výjimku zachytit v rámci __exit__ funkce, nějakým způsobem si s ní poradit a nepropagovat ji ven. Pokud by funkce __exit__ vracela True, znamenalo by to, že se s výjimkou vypořádal kontextový manažer a už není potřeba ji řešit o úroveň výš.

Kontextový manažer pomocí generátorů

Dalším skvělým způsobem, jak si "zapamatovat" vnitřní stav a definovat si tak nějaký kontext, je využít generátory. Generátory jsou, stručně řečeno, funkce, které se dají "pozastavit" a při opětovném zavolání "spustit" od tohoto místa dál. Tady vidíme skvělou příležitost pro kontextové manažery, protože jsme schopni v první části takové funkce nainicializovat všechny potřebné proměnné, pomocí klíčového slova yield vrátit, co je potřeba, hlavnímu procesu a po ukončení with bloku pokračovat ve funkci dál. Tento zápis si rovněž bez problému poradí se zachytáváním výjimek.

from contextlib import contextmanager

@contextmanager
def file_manager(path, mode):
    file = open(file_path, mode)
    try:
        yield file
    finally:
        file.close()


with file_manager("soubor.txt", "w") as f:
    f.write("Ahoj světe!")