Context Managers (with)

Difficulty: Intermediate

Question

What are context managers? How do you create custom ones?

Answer

Context managers ensure resources are properly acquired and released using the 'with' statement. They implement the context manager protocol: __enter__ and __exit__.

__enter__: Called when entering the with block, returns the resource. __exit__: Called when exiting (even if an exception occurred), handles cleanup.

Common uses: file handling, database connections, locks, temporary changes, timing.

Two ways to create context managers: 1. Class-based: implement __enter__ and __exit__ 2. Function-based: use @contextlib.contextmanager decorator with yield

Code examples

Class-based Context Manager

class Timer:
    def __init__(self, label="Operation"):
        self.label = label
    
    def __enter__(self):
        import time
        self.start = time.perf_counter()
        return self  # Value assigned to 'as' variable
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        import time
        self.elapsed = time.perf_counter() - self.start
        print(f"{self.label} took {self.elapsed:.4f}s")
        return False  # Don't suppress exceptions

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

# Exception handling in __exit__
class SuppressErrors:
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is ValueError:
            print(f"Suppressed ValueError: {exc_val}")
            return True  # True = suppress the exception
        return False     # False = re-raise

with SuppressErrors():
    raise ValueError("test error")  # Suppressed!
print("Continues after suppressed error")

__exit__ receives exception info (or None if no exception). Returning True suppresses the exception, False (or None) re-raises it.

Function-based with @contextmanager

from contextlib import contextmanager

@contextmanager
def temp_directory():
    import tempfile, shutil, os
    dirpath = tempfile.mkdtemp()
    print(f"Created temp dir: {dirpath}")
    try:
        yield dirpath  # This is what 'as' receives
    finally:
        shutil.rmtree(dirpath)
        print(f"Cleaned up temp dir")

with temp_directory() as tmpdir:
    print(f"Working in: {tmpdir}")
    # Create files, do work...
# Automatically cleaned up after the block

@contextmanager
def database_transaction(db):
    """Commit on success, rollback on failure."""
    try:
        yield db
        db.commit()
        print("Transaction committed")
    except Exception:
        db.rollback()
        print("Transaction rolled back")
        raise

# Simulating
class FakeDB:
    def commit(self): pass
    def rollback(self): pass

with database_transaction(FakeDB()) as db:
    print("Doing database work...")

@contextmanager turns a generator function into a context manager. Code before yield is __enter__, code after yield is __exit__. The yield value is the 'as' target.

Built-in Context Managers

# File handling (most common)
with open('example.txt', 'w') as f:
    f.write('Hello, World!')
# File is automatically closed, even if an exception occurs

# Multiple context managers
with open('input.txt') as src, open('output.txt', 'w') as dst:
    dst.write(src.read())

# Threading lock
import threading
lock = threading.Lock()
with lock:  # Acquires lock, releases on exit
    print("Critical section")

# Suppress specific exceptions
from contextlib import suppress
with suppress(FileNotFoundError):
    import os
    os.remove('nonexistent.txt')  # No error raised
print("Continues normally")

# Redirect stdout
from contextlib import redirect_stdout
import io
f = io.StringIO()
with redirect_stdout(f):
    print("This goes to the buffer")
captured = f.getvalue()
print(f"Captured: {captured.strip()}")

Python's standard library is full of context managers: files, locks, database connections, temporary directories, output redirection, and exception suppression.

Key points

• Always use 'with' for file operations - it guarantees the file is closed even on exceptions
• contextlib.suppress() is cleaner than try/except/pass for ignoring specific exceptions
• contextlib.ExitStack() manages a dynamic number of context managers

Concepts covered

with statement, __enter__, __exit__, contextlib, Resource Management