context_manager_demo.py

script

← Back to skill

Content hash: 36bad01c86feb666ef42f65bf828445916a718d974cf90498a8653b3d5d64cb6
#!/usr/bin/env python3
"""Context manager patterns: class-based, generator-based, ExitStack, and pitfalls.

Covers: __enter__/__exit__, @contextmanager, ExitStack for dynamic resources,
and common bugs (swallowing exceptions, releasing unacquired resources).
"""
from __future__ import annotations

import contextlib
import sqlite3
import tempfile
from pathlib import Path
from threading import Lock


# ---- 1. Class-based: resource acquisition + guaranteed cleanup ----

class DatabaseConnection:
    def __init__(self, path: str):
        self._path = path
        self._conn: sqlite3.Connection | None = None

    def __enter__(self) -> sqlite3.Connection:
        self._conn = sqlite3.connect(self._path)
        self._conn.execute("PRAGMA journal_mode=WAL")
        return self._conn

    def __exit__(self, exc_type, exc_val, exc_tb) -> bool | None:
        if self._conn:
            self._conn.close()
        return None  # don't suppress exceptions


# ---- 2. Generator-based: simpler for read-only or quick patterns ----

@contextlib.contextmanager
def temp_file(content: str, suffix: str = ".txt"):
    path = None
    try:
        with tempfile.NamedTemporaryFile(mode="w", suffix=suffix, delete=False) as f:
            f.write(content)
            path = Path(f.name)
        yield path
    finally:
        if path and path.exists():
            path.unlink()


# ---- 3. ExitStack: dynamic number of resources ----

def open_many_files(paths: list[str]) -> list[str]:
    with contextlib.ExitStack() as stack:
        files = [stack.enter_context(open(p)) for p in paths]
        return [f.readline().strip() for f in files]


# ---- 4. Thread-safe lock ----

class LockedCounter:
    def __init__(self):
        self._lock = Lock()
        self._value = 0

    @contextlib.contextmanager
    def atomic_inc(self):
        with self._lock:
            self._value += 1
            yield self._value


# ---- Demo ----

def main() -> None:
    # Class-based
    with DatabaseConnection(":memory:") as conn:
        conn.execute("CREATE TABLE test (id int)")
        conn.execute("INSERT INTO test VALUES (1)")
        row = conn.execute("SELECT * FROM test").fetchone()
        print(f"DB row: {row}")

    # Generator-based
    with temp_file("hello world") as path:
        content = path.read_text()
        print(f"Temp file content: {content}")
    print(f"Temp file cleaned up: {not path.exists()}")

    # ExitStack
    p1 = Path(tempfile.mktemp())
    p1.write_text("line1")
    p2 = Path(tempfile.mktemp())
    p2.write_text("line2")
    lines = open_many_files([str(p1), str(p2)])
    print(f"Lines: {lines}")
    p1.unlink()
    p2.unlink()

    # Locked counter
    counter = LockedCounter()
    for i in range(3):
        with counter.atomic_inc() as val:
            print(f"Counter: {val}")

    print("\nAll context managers worked correctly.")


if __name__ == "__main__":
    main()