pythonintermediate

Python Context Manager Patterns

Create custom context managers with __enter__/__exit__ and contextlib for resource management.

python
from contextlib import contextmanager, asynccontextmanager
import time
import tempfile
import os


# Class-based context manager
class Timer:
    def __enter__(self):
        self.start = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.elapsed = time.perf_counter() - self.start
        print(f"Elapsed: {self.elapsed:.4f}s")
        return False  # Don't suppress exceptions


# Generator-based context manager
@contextmanager
def temp_directory():
    dirpath = tempfile.mkdtemp()
    try:
        yield dirpath
    finally:
        import shutil
        shutil.rmtree(dirpath)


# Database transaction context manager
@contextmanager
def transaction(conn):
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise


# Change directory temporarily
@contextmanager
def cd(path: str):
    old = os.getcwd()
    os.chdir(path)
    try:
        yield
    finally:
        os.chdir(old)


# Suppress specific exceptions
@contextmanager
def suppress(*exceptions):
    try:
        yield
    except exceptions:
        pass


# Usage
with Timer() as t:
    total = sum(range(1_000_000))
print(f"Sum: {total}, Time: {t.elapsed:.4f}s")

with temp_directory() as d:
    filepath = os.path.join(d, "test.txt")
    with open(filepath, "w") as f:
        f.write("temporary data")
# Directory is automatically cleaned up

Use Cases

  • Automatic resource cleanup and lifecycle management
  • Database transaction handling
  • Temporary environment modifications

Tags

Related Snippets

Similar patterns you can reuse in the same workflow.