Python Context Managers: Beyond with open()

#python #tutorial #programming #bestpractices

Python Context Managers: Beyond with open()

Context managers are not just for files. Here's how to create your own.

The Protocol

class ManagedResource:
    def __enter__(self):
        print("Acquiring resource")
        return self  # returned to the 'as' variable

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("Releasing resource")
        return False  # Don't suppress exceptions

with ManagedResource() as res:
    print("Using resource")
# Output:
# Acquiring resource
# Using resource
# Releasing resource

contextlib.contextmanager

from contextlib import contextmanager
import time

@contextmanager
def timer(name: str = ""):
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start
        print(f"{name or 'Block'} took {elapsed:.4f}s")

with timer("database query"):
    time.sleep(0.1)
# database query took 0.1001s

Database Transaction Manager

from contextlib import contextmanager
import sqlite3

@contextmanager
def transaction(db_path: str):
    conn = sqlite3.connect(db_path)
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()

with transaction("app.db") as conn:
    conn.execute("INSERT INTO users VALUES (?, ?)", (1, "Alice"))
# Auto-commits on success, rollbacks on error

Suppress Exceptions

from contextlib import suppress

# Instead of:
try:
    os.remove("temp.txt")
except FileNotFoundError:
    pass

# Use:
with suppress(FileNotFoundError):
    os.remove("temp.txt")

Async Context Managers

from contextlib import asynccontextmanager
import httpx

@asynccontextmanager
async def http_client(base_url: str):
    async with httpx.AsyncClient(base_url=base_url) as client:
        yield client

async def main():
    async with http_client("https://api.example.com") as client:
        r = await client.get("/users")
        print(r.json())