Context Managers - Python Mastery

Context Managers: Resource Management

Context managers handle setup and cleanup automatically. The with statement calls __enter__ on entry and __exit__ on exit — even if an exception occurs. You've already used them with files.

# File I/O is the classic context manager
with open("file.txt", "w") as f:
    f.write("content")
# File is automatically closed here, even if an exception occurred

# What it's equivalent to:
f = open("file.txt", "w")
try:
    f.write("content")
finally:
    f.close()

# Multiple context managers in one with statement
with open("input.txt") as fin, open("output.txt", "w") as fout:
    for line in fin:
        fout.write(line.upper())

Building Context Managers with Classes

class Timer:
    """Time a block of code."""
    import time as _time

    def __enter__(self):
        self.start = __import__('time').perf_counter()
        return self    # Bound to 'as' variable

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

with Timer() as t:
    # Do some work
    sum(range(1_000_000))
# Elapsed: 0.0321s

# Database connection example
class DatabaseConnection:
    def __init__(self, db_url):
        self.db_url = db_url
        self.conn = None

    def __enter__(self):
        self.conn = connect(self.db_url)
        return self.conn

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type:
            self.conn.rollback()   # Error: roll back
        else:
            self.conn.commit()     # Success: commit
        self.conn.close()
        return False   # Let exceptions propagate

with DatabaseConnection("sqlite:///mydb.db") as db:
    db.execute("INSERT INTO users VALUES (?)", ("Alice",))

contextlib: Easier Context Managers

from contextlib import contextmanager

# Use a generator instead of a class
@contextmanager
def timer():
    import time
    start = time.perf_counter()
    yield    # Code in 'with' block runs here
    elapsed = time.perf_counter() - start
    print(f"Elapsed: {elapsed:.4f}s")

with timer():
    sum(range(1_000_000))

# Temporary directory
@contextmanager
def temp_directory():
    import tempfile, shutil
    tmpdir = tempfile.mkdtemp()
    try:
        yield tmpdir
    finally:
        shutil.rmtree(tmpdir)   # Always clean up

with temp_directory() as tmpdir:
    # Work in tmpdir — it's deleted after the block
    with open(f"{tmpdir}/test.txt", "w") as f:
        f.write("temporary data")

# contextlib.suppress — suppress specific exceptions
from contextlib import suppress

with suppress(FileNotFoundError):
    os.remove("might_not_exist.txt")   # No error if missing

Key Takeaways

with statement = guaranteed cleanup even when exceptions occur
__enter__ sets up, __exit__ tears down
@contextmanager: the generator approach is simpler than a class
yield splits the context manager: code before = __enter__, after = __exit__
contextlib.suppress: clean way to ignore specific exceptions

Practice Exercises

Write a @contextmanager called redirect_stdout that captures print output to a string.
Write a context manager that changes the working directory and restores it afterwards.
Write a context manager for a "transaction": on exit, call commit() if no exception, rollback() if there was one.

← Exceptions Modules and Imports →