Context Manager ochroni Twoje plecy!

Menadżer kontekstu (ang. context manager) to obiekt, który odpowiednio zarządza danym zasobem, zapewniając, że zostanie on odpowiednio zamknięty. Przez zamknięty mam tu na myśli czyszczenie, zwalnianie zasobów oraz sprzątanie po wykonaniu bloku kodu. Dobrym przykładem takiego obiektu, jest użycie open w Pythonie. Korzystając z niego w domyślny sposób, otwieramy plik, odczytujemy pewien jego element, a następnie go zamykamy.

1
f = open('versions.txt')
2
print(f.readline())
3
f.close()

Jednak open to też menadżer kontekstu, aby z niego skorzystać, należy skorzystać z konstrukcji with ... as ... :, gdzie samo as nie jest wymagane gdy nie potrzebujemy w danym kontekście korzystać z jego obiektu. Więc po with podajemy nasz menadżer, czyli open(). Następnie, aby mieć dostęp do otwieranego pliku, dodajemy as i zmienną, ja użyję zmiennej f.

1
with open('base.py') as f:
2
    print(f.readline())

Dzięki ten konstrukcji nie musimy korzystać z f.close(), plik zostanie zamknięty samoczynnie po wyjściu z kontekstu, czyli po zakończeniu działania na kodzie, który znalazł się w jej ciele poprzez wcięcie.

Baza

Aby pokazać więcej na ten temat, utworzyłem klasę Storage, która będzie symulować magazyn danych. Dla zabawy, przy próbie pobrania czegokolwiek z zasobu możemy dostać wyjątek StorageException. Aby z niego korzystać, musimy uprzednio zablokować dostęp do zasobu, a na koniec zasób zwolnić.

1
import random
2

3
class StorageException(Exception):
4
    """Raised when there is some issue with storage"""
5

6
class Storage():
7

8
    def __init__(self):
9
        print('init')
10

11
    def block_resources(self):
12
        print('block_resources')
13

14
    def release_resources(self):
15
        print("release_resources")
16

17
    def get(self, *args, **kwargs):
18
        if random.randint(0, 1):
19
            raise StorageException("Something goes wrong")
20
        print("get", args, kwargs)

Domyślne użycie tej klasy wyglądałoby następująco.

1
from base import Storage
2
storage = Storage()
3
storage.block_resources()
4
storage.get(filename='secrets.yaml')
5
storage.release_resources()

1
init
2
block_resources
3
get () {'filename': 'secrets.yaml'}
4
release_resources

O ile nie wystąpiłby błąd, wszystko byłoby ok. Jednak czasem przy wywołaniu możemy dostać błąd. Spowoduje to nie wykonanie się ostatniej linii kodu, a to oznacza zablokowanie zasobu do momentu, gdy ktoś manualnie go odblokuje lub proces utrzymujący dostęp do zasobu zostanie zabity i system przywróci zasób do swojej puli.

1
init
2
block_resources
3
Traceback (most recent call last):
4
    storage.get(filename='secrets.yaml')
5
    raise StorageException("Something goes wrong")
6
base.StorageException: Something goes wrong

Dispose Pattern

Taki problem rozwiązuje wzorzec projektowy Dispose, który służy właśnie do zamykania zasobów w przypadku gdy pojawi się błąd. Jego najprostsza implementacja w naszym przypadku korzysta z try finally, więc jeżeli cokolwiek wydarzy się w naszym kodzie, to przed ostatecznym wywołaniem wyjątku i zakończeniem działania kodu wykona się kod sekcji finally. Dzięki temu mamy pewność, że nasz zasób zawsze zostanie zamknięty, czy też w tym przypadku zwolniony.

1
from base import Storage
2

3
storage = Storage()
4
storage.block_resources()
5

6
try:
7
    storage.get(filename='secrets.yaml')
8
finally:
9
    storage.release_resources()

Mała uwaga, nie próbuj w bloku finally zwracać żadnej wartości. Spowoduje to nadpisanie tego, co powinno zostać zwrócone. Utracisz w ten sposób wartość zwracaną z kodu, lub zignorujesz wyjątek.

Elementy menadżera kontekstu

Przejdźmy do stworzenia własnego menadżera, bez korzystania z dokumentacji. Będziemy zatem obserwować błędy wyrzucane po uruchomieniu kodu, i dodawać do niego brakujące elementy. Więc w pierwszej kolejności tworzymy pustą klasę, i spróbujemy ją użyć z instrukcją with.

1
class StorageContext:
2
    pass
3

4
with StorageContext() as storage:
5
    print(storage)

1
    with StorageContext() as storage:
2
AttributeError: __enter__

Wyjątek, który się pojawił, mówi jasno, iż brakuje nam atrybutu enter. Podejrzewam, że chodzi o metodę, dlatego dodaje do naszej klasy taką metodę. Następnie uruchamiam kod ponownie.

1
class StorageContext:
2

3
    def __enter__(self):
4
        pass
5

6
with StorageContext() as storage:
7
    print(storage)

1
    with StorageContext() as storage:
2
AttributeError: __exit__

To ma sens. Mamy metodę enter do otwierania pewnego kontekstu, ale potrzebujemy także exit aby ten kontekst zakończyć. Dodajemy metodę i uruchamiamy kod ponownie.

1
class StorageContext:
2

3
    def __enter__(self):
4
        pass
5

6
    def __exit__(self):
7
        pass
8

9
with StorageContext() as storage:
10
    print(storage)

1
TypeError: __exit__() takes 1 positional argument but 4 were given

Dostajemy informacje, że metoda exit() przyjmuje aż 4 argumenty pozycyjne. Jeden to self, resztę opakujemy w *args, sprawdzimy potem, co się w nich kryje. Uruchamiamy ponownie.

1
class StorageContext:
2

3
    def __enter__(self):
4
        pass
5

6
    def __exit__(self, *args):
7
        pass
8

9
with StorageContext() as storage:
10
    print(storage)

1
None

Kod uruchomił się bez błędów. Udało się zdefiniować szkielet, jakiego wymaga Python. Popracujmy dalej nad tym, aby funkcja print zwróciła None. Podejrzewam, że metoda enter powinna zwrócić jakiś obiekt, na którym będziemy pracować, tak więc dodaje return self.

1
class StorageContext:
2

3
    def __enter__(self):
4
        return self
5

6
    def __exit__(self, *args):
7
        pass
8

9
with StorageContext() as storage:
10
    print(storage)

1
<__main__.StorageContext object at 0x7fc769596ee0>

Tak jak podejrzewałem, aby móc pracować na obiekcie ograniczonym naszym kontekstem, w metodzie enter należy ten obiekt zwrócić. Została ostatnia magiczna rzecz, a jest nią zbiór argumentów pozycyjnych w metodzie exit. Dodajmy tam print aby się przekonać.

1
class StorageContext:
2

3
    def __enter__(self):
4
        return self
5

6
    def __exit__(self, *args):
7
        print(args)
8

9
with StorageContext() as storage:
10
    print(storage)

1
<__main__.StorageContext object at 0x7f1f9aaf2ee0>
2
(None, None, None)

Jest pusto, to dziwne. Jednak pamiętajmy, że wzorzec projektowy korzystał z obsługi błędów, więc może jego wywołanie coś zmieni.

1
class StorageContext:
2

3
    def __enter__(self):
4
        return self
5

6
    def __exit__(self, *args):
7
        print(args)
8

9
with StorageContext() as storage:
10
    raise ValueError

1
(<class 'ValueError'>, ValueError(), <traceback object at 0x10ab4ae88>)
2
Traceback (most recent call last):
3
    raise ValueError
4
ValueError

Mamy wyjaśnienie. Metoda exit dostaje informacje o wyjątkach, które się pojawiły. Dzięki temu możemy te wyjątki odpowiednio obsłużyć lub ich obsługę pominąć. Jest jeszcze jedna rzecz, którą ciężko by mi było wypracować. W momencie gdy exit zwraca wartość domyślną None, to wyjątek i tak zostanie wyrzucony. Aby do tego nie dopuścić, należy zwrócić z metody exit wartość True.

Implementacja za pomocą klasy

Ostateczna implementacja naszego kodu do obsługi zasobu za pomocą klasy wygląda zatem następująco:

1
from base import Storage
2

3
class StorageContext:
4

5
    def __init__(self):
6
        self.storage = Storage()
7

8
    def __enter__(self):
9
        self.storage.block_resources()
10
        return self.storage
11

12
    def __exit__(self, exc_type, exc_val, exc_tb):
13
        self.storage.release_resources()
14

15
with StorageContext() as s:
16
    s.get(filename='secrets.yaml')

Implementacja za pomocą funkcji

Ktoś mądry pomyślał, że dobrym pomysłem będzie tworzenie prostych menadżerów kontekstu opartych o generator. W taki sposób powstał dekorator contextmanager z biblioteki wbudowanej contextlib, dzięki któremu możemy przekształcić generator w menadżera kontekstu. Sposób implementacji naszego problemu za pomocą tego narzędzia poniżej. Rozwiązanie to jest bardzo często używane, gdyż ma prostszą konstrukcję oraz dużo łatwiej go zaimplementować.

1
from contextlib import contextmanager
2
from base import Storage
3

4
@contextmanager
5
def storage_access():
6
    storage = Storage()
7
    storage.block_resources()
8
    try:
9
        yield storage
10
    finally:
11
        storage.release_resources()
12

13
with storage_access() as s:
14
    s.get(filename='secrets.yaml')

Przykłady

Skoro wiemy już, jak tworzyć context manager, to co można z nim zrobić? Jeden z najważniejszych przypadków pokazałem na powyższych przykładach, czyli dostęp do pewnych zasobów, które wymagają konkretnej operacji przed i po.

Code Timer

Inny ciekawy przykład, to stworzenie klasy do mierzenia czasu działania danego kawałka kodu. Z reguły jest to przykład na dekorator, ale wymaga on wtedy zamknięcia naszego kodu w funkcji. Licznik oparty o menadżera kontekstu nie ma tego ograniczenia.

Dodatkowo pojawia się tutaj prosta konfiguracja logera, dzięki której możemy logować informacje z naszego kodu. Korzystanie z logera jest dużo lepszym wzorcem niż korzystanie z print w naszym kodzie, ponieważ loger możemy odpowiednio skonfigurować, aby dla przykładu wysyłał nasze logi do platformy , gdzie te logi możemy następnie przejrzeć.

1
import time
2
import logging
3

4
logging.basicConfig(level=logging.DEBUG)
5

6
class LogElapsedTime:
7
    def __init__(self, name=""):
8
        self.start = None
9
        self.name = name
10

11
    def __enter__(self):
12
        self.start = time.time()
13

14
    def __exit__(self, *args):
15
        end = time.time()
16
        delta = end - self.start
17
        logging.info("time of execution '%s':%s", self.name, delta)
18

19
with LogElapsedTime("sleep"):
20
    time.sleep(1)

1
INFO:root:time of execution 'sleep':1.0010440349578857

Ignore Exceptions

Inną propozycją ode mnie, jest context manager, którego zadaniem będzie ignorowanie wszelkich błędów, które wystąpią. Nie chcemy jednak stracić informacji o tym, iż wystąpiły, dlatego odpowiednia informacja zostanie przekazana do logera.

1
import logging
2
logging.basicConfig(level=logging.DEBUG)
3
class IgnoreExceptions:
4
    def __enter__(self):
5
        pass
6
    def __exit__(self, exc_type, exc_val, exc_tb):
7
        if exc_type:
8
            logging.error("exception ignored : %s", exc_type.__name__)
9
        return True
10
with IgnoreExceptions():
11
    raise ArithmeticError

1
ERROR:root:exception ignored : ArithmeticError