Menedżery kontekstu w Pythonie: instrukcja with wyjaśniona

Instrukcja, która sprząta po sobie

Każdy zasób otwarty w programie (plik, połączenie sieciowe, uchwyt do bazy danych, blokada) trzeba zamknąć, gdy skończysz. Zapomnisz, a pojawią się wycieki pamięci, blokady wstrzymujące inne procesy albo uszkodzone pliki po awarii. Instrukcja with w Pythonie załatwia to za ciebie.

Wzorzec, od którego zaczyna każdy, to odczyt pliku:

with open("notes.txt") as f:
    contents = f.read()
    print(contents)

Dwie rzeczy dzieją się automatycznie. Przy wejściu open() daje ci obiekt pliku związany z f. Przy wyjściu, niezależnie od tego, czy blok kończy się normalnie, wcześniej wraca, czy zgłasza wyjątek, Python wywołuje za ciebie f.close().

To wszystko. Na tym polega cały sens with.

Co zastępuje "with"

Przed menedżerami kontekstu bezpieczny odpowiednik wymagał try/finally:

f = open("notes.txt")
try:
    contents = f.read()
    print(contents)
finally:
    f.close()

Pięć linii ceremonii tylko po to, by "odczytać plik i zamknąć go na koniec". Pomnóż to przez każde open w większym programie, a zalety staną się oczywiste. with jest krótsze, trudniej w nim o błąd i nie da się zapomnieć o sprzątaniu.

Otwieranie kilku zasobów

W jednym with możesz powiązać kilka menedżerów kontekstu:

with open("input.txt") as src, open("output.txt", "w") as dst:
    dst.write(src.read().upper())

Oba pliki otwierają się przy wejściu i oba zamykają przy wyjściu. Jeśli pierwsze open się uda, a drugie zgłosi wyjątek, Python i tak zamknie pierwszy plik: mechanizm poprawnie obsługuje częściowe przygotowanie.

Przy dłuższych listach zasobów czytelniejsza jest forma z nawiasami (Python 3.10+):

with (
    open("a.txt") as a,
    open("b.txt") as b,
    open("c.txt") as c,
):
    ...

Czym naprawdę jest menedżer kontekstu

Każdy obiekt, który definiuje __enter__ i __exit__, jest menedżerem kontekstu. Protokół jest bardzo prosty:

• __enter__(self) uruchamia się, gdy zaczyna się blok with. Jego wartość zwracana trafia do nazwy po as name.
• __exit__(self, exc_type, exc_value, traceback) uruchamia się, gdy blok się kończy, niezależnie od tego, w jaki sposób. Jeśli wyjście spowodował wyjątek, informacje o nim są przekazywane, żeby menedżer kontekstu mógł go zbadać albo stłumić.

Oto minimalny menedżer, który mierzy czas wykonania bloku:

import time

class Timer:
    def __enter__(self):
        self.start = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        elapsed = time.perf_counter() - self.start
        print(f"Elapsed: {elapsed*1000:.2f} ms")

with Timer():
    total = sum(i * i for i in range(1_000_000))
    print(total)

with Timer(): tworzy obiekt, wywołuje jego __enter__, wykonuje ciało bloku i wywołuje __exit__. Bez pliku, bez blokady: tylko mała otoczka wokół "zrób coś i zmierz, ile to trwało".

Skrót contextlib.contextmanager

Definiowanie klasy dla każdego menedżera kontekstu to więcej pracy, niż trzeba. contextlib.contextmanager zamienia funkcję generatora w menedżer kontekstu, a jedno yield oddziela "przed" od "po":

from contextlib import contextmanager
import time

@contextmanager
def timed(label):
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start
        print(f"{label}: {elapsed*1000:.2f} ms")

with timed("big sum"):
    total = sum(i * i for i in range(1_000_000))
    print(total)

Wszystko przed yield to zachowanie __enter__. Wszystko po nim to __exit__. Dzięki try/finally sprzątanie wykona się nawet wtedy, gdy ciało bloku zgłosi wyjątek.

Większość własnych menedżerów kontekstu, które napiszesz, pasuje do tego kształtu. Najpierw sięgaj po formę z dekoratorem, a klasę pisz tylko wtedy, gdy potrzebujesz czegoś, czego forma z generatorem nie wyrazi.

Tymczasowa zmiana czegoś

Częsty schemat: ustaw coś, użyj, przywróć. Menedżery kontekstu wyrażają to czysto:

from contextlib import contextmanager

@contextmanager
def cd(path):
    import os
    previous = os.getcwd()
    os.chdir(path)
    try:
        yield
    finally:
        os.chdir(previous)

# Uruchamia kod w innym katalogu roboczym,
# a po zakończeniu bloku automatycznie wraca.
with cd("/tmp"):
    import os
    print("cwd:", os.getcwd())

import os
print("back to:", os.getcwd())

Każdy schemat "ustaw, a potem przywróć" (zmienne środowiskowe, poziom szczegółowości logów, flagi funkcji, fixture'y testów) naturalnie mieści się w menedżerze kontekstu. Wywołujący nie musi pamiętać o przywracaniu czegokolwiek.

Tłumienie wyjątków

Metoda __exit__ może zwrócić True, by powiedzieć Pythonowi "wyjątek obsłużony, połknij go". To rzadkie i zwykle podejrzane, ale tak działa contextlib.suppress:

from contextlib import suppress

with suppress(FileNotFoundError):
    with open("maybe.txt") as f:
        print(f.read())

print("carried on")

suppress(FileNotFoundError) zamienia FileNotFoundError w operację, która nic nie robi. Używaj tego do naprawdę opcjonalnych operacji w stylu "spróbuj, a jak nie wyjdzie, trudno". Nie używaj go do uciszania wyjątków, których nie przemyślano.

Inne menedżery kontekstu, które spotkasz

Gdy zaczniesz się rozglądać, menedżery kontekstu znajdziesz w całej bibliotece standardowej:

import threading
from pathlib import Path

# Blokady: gwarantują zwolnienie, nawet gdy sekcja krytyczna zgłosi wyjątek.
lock = threading.Lock()
with lock:
    ...

# tempfile: usuwa plik tymczasowy, gdy skończysz.
from tempfile import TemporaryDirectory
with TemporaryDirectory() as tmp:
    path = Path(tmp) / "scratch.txt"
    path.write_text("hello")

# Połączenia z bazą danych: zamykają połączenie (albo kończą transakcję).
import sqlite3
with sqlite3.connect(":memory:") as conn:
    conn.execute("CREATE TABLE t (x INTEGER)")

Biblioteki zewnętrzne trzymają się tych samych konwencji. Gdy widzisz with something as x:, prawie zawsze znaczy to "używaj x przez czas trwania tego bloku, a potem posprzątaj".

Kiedy nie używać with

• Gdy tak naprawdę nie ma przygotowania ani sprzątania. Opakowywanie dowolnego kodu w menedżer kontekstu bez powodu tylko dodaje szumu.
• Gdy zasób jest potrzebny w wielu niezwiązanych blokach. Trzymanie otwartego with przez cały czas działania długiego skryptu może ukryć, jaki jest naprawdę zakres sprzątania. Zamiast tego rozważ klasę, która jest właścicielem zasobu.
• Gdy lepiej pasuje dekorator. Niektóre powtarzalne wzorce (ponawianie, logowanie, pomiar czasu) czyta się naturalniej jako @decorator na funkcji niż jako with ...: w jej wnętrzu. Wybierz to, co czyta się lepiej w miejscu wywołania.

W większości przypadków with to dobry wybór. Rzadkie wyjątki łatwo zauważyć, gdy już wiesz, czego szukać.

Dalej: praca z prawdziwymi plikami

Znasz już mechanizm stojący za with open(...) as f:, czyli kontekstem, w którym użyjesz go w dziewięćdziesięciu procentach przypadków. Następny rozdział wykorzystuje go do odczytu, zapisu i poruszania się po plikach na dysku.