Context Managers

In this chapter, we'll learn about context managers - Python's elegant way to manage resources. We'll cover:

• What are context managers?
• The with statement
• Creating context managers with __enter__ and __exit__
• Using the @contextmanager decorator
• File handling and other common use cases
• Testing context managers

What is a context manager?

A context manager handles setup and cleanup automatically. The most common example is file handling:

# Without context manager - you must remember to close
file = open("data.txt", "r")
try:
    content = file.read()
finally:
    file.close()

# With context manager - cleanup is automatic
with open("data.txt", "r") as file:
    content = file.read()
# File is automatically closed here, even if an exception occurs

The with statement

The with statement:

1. Calls __enter__ at the start
2. Runs your code block
3. Calls __exit__ when the block ends (even if an exception occurs)

with resource as value:
    # Use value
# Cleanup happens automatically

Creating a context manager (class-based)

Write the test first

Let's create a timer context manager:

import time
from timer import Timer


def test_timer_measures_duration():
    with Timer() as timer:
        time.sleep(0.1)
    
    assert timer.duration >= 0.1
    assert timer.duration < 0.2


def test_timer_works_with_exception():
    try:
        with Timer() as timer:
            time.sleep(0.05)
            raise ValueError("test")
    except ValueError:
        pass
    
    assert timer.duration >= 0.05

Implementation

import time


class Timer:
    def __init__(self):
        self.start_time = None
        self.end_time = None
        self.duration = None
    
    def __enter__(self):
        self.start_time = time.time()
        return self  # This becomes the 'as' value
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end_time = time.time()
        self.duration = self.end_time - self.start_time
        return False  # Don't suppress exceptions

Understanding exit parameters

def __exit__(self, exc_type, exc_val, exc_tb):
    # exc_type: Exception class (or None if no exception)
    # exc_val: Exception instance (or None)
    # exc_tb: Traceback (or None)
    
    # Return True to suppress the exception
    # Return False (or None) to propagate it
    return False

The @contextmanager decorator

A simpler way to create context managers using generators:

from contextlib import contextmanager


@contextmanager
def timer():
    start = time.time()
    
    class Result:
        duration = None
    
    result = Result()
    
    try:
        yield result  # Everything before yield is __enter__
    finally:
        result.duration = time.time() - start  # This is __exit__

Testing the generator-based version

from timer import timer


def test_timer_context_manager():
    with timer() as t:
        time.sleep(0.1)
    
    assert t.duration >= 0.1

Common use cases

Database connections

@contextmanager
def database_connection(connection_string):
    conn = connect(connection_string)
    try:
        yield conn
    finally:
        conn.close()


# Usage
with database_connection("postgresql://...") as conn:
    conn.execute("SELECT * FROM users")

Temporary file cleanup

import os
from contextlib import contextmanager


@contextmanager
def temp_file(content):
    filename = f"/tmp/temp_{os.getpid()}.txt"
    with open(filename, "w") as f:
        f.write(content)
    try:
        yield filename
    finally:
        os.remove(filename)


# Usage
with temp_file("Hello, World!") as path:
    with open(path, "r") as f:
        print(f.read())
# File is deleted after the block

Changing directories

import os
from contextlib import contextmanager


@contextmanager
def change_directory(path):
    original = os.getcwd()
    try:
        os.chdir(path)
        yield
    finally:
        os.chdir(original)


# Usage
with change_directory("/tmp"):
    # Work in /tmp
    print(os.getcwd())
# Back to original directory

Mocking for tests

@contextmanager
def mock_environment(**env_vars):
    original = os.environ.copy()
    os.environ.update(env_vars)
    try:
        yield
    finally:
        os.environ.clear()
        os.environ.update(original)


# Usage
def test_with_mock_env():
    with mock_environment(API_KEY="test123"):
        assert os.environ["API_KEY"] == "test123"

Nested context managers

Python allows multiple context managers in one with:

# Multiple files
with open("input.txt") as infile, open("output.txt", "w") as outfile:
    outfile.write(infile.read())

# Equivalent to nested with statements
with open("input.txt") as infile:
    with open("output.txt", "w") as outfile:
        outfile.write(infile.read())

Exception handling in context managers

Suppressing exceptions

@contextmanager
def suppress_exceptions(*exception_types):
    try:
        yield
    except exception_types:
        pass  # Suppress specified exceptions


# Usage
with suppress_exceptions(ValueError, TypeError):
    raise ValueError("This is suppressed")

print("Continues normally")

Python's contextlib provides this:

from contextlib import suppress

with suppress(ValueError):
    raise ValueError("Suppressed")

Logging exceptions

@contextmanager
def log_exceptions(logger):
    try:
        yield
    except Exception as e:
        logger.error(f"Exception occurred: {e}")
        raise  # Re-raise the exception

Testing context managers

Testing enter return value

def test_database_connection_returns_connection():
    with DatabaseConnection("test://") as conn:
        assert conn is not None
        assert conn.is_connected()

Testing cleanup happens

def test_temp_file_is_deleted():
    with temp_file("test content") as path:
        assert os.path.exists(path)
    
    assert not os.path.exists(path)

Testing exception handling

def test_cleanup_on_exception():
    resource = MockResource()
    
    try:
        with resource:
            raise ValueError("test")
    except ValueError:
        pass
    
    assert resource.was_cleaned_up

A complete example: Database transaction

# test_transaction.py
import pytest
from database import Transaction, DatabaseError


def test_transaction_commits_on_success():
    mock_db = MockDatabase()
    
    with Transaction(mock_db) as tx:
        tx.execute("INSERT INTO users VALUES (1, 'Alice')")
    
    assert mock_db.committed
    assert not mock_db.rolled_back


def test_transaction_rollback_on_error():
    mock_db = MockDatabase()
    
    with pytest.raises(ValueError):
        with Transaction(mock_db) as tx:
            tx.execute("INSERT INTO users VALUES (1, 'Alice')")
            raise ValueError("Something went wrong")
    
    assert not mock_db.committed
    assert mock_db.rolled_back

Implementation:

class Transaction:
    def __init__(self, database):
        self.db = database
        self.connection = None
    
    def __enter__(self):
        self.connection = self.db.get_connection()
        self.connection.begin()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            self.connection.commit()
        else:
            self.connection.rollback()
        self.connection.close()
        return False  # Don't suppress exceptions
    
    def execute(self, query):
        self.connection.execute(query)

contextlib utilities

The contextlib module provides helpful utilities:

from contextlib import (
    contextmanager,      # Decorator for generator-based context managers
    suppress,            # Suppress specified exceptions
    redirect_stdout,     # Redirect stdout
    redirect_stderr,     # Redirect stderr
    nullcontext,         # No-op context manager
    closing,             # Call close() on exit
    ExitStack,           # Manage multiple context managers dynamically
)

ExitStack for dynamic contexts

from contextlib import ExitStack


def process_files(filenames):
    with ExitStack() as stack:
        files = [
            stack.enter_context(open(fn))
            for fn in filenames
        ]
        # All files are open here
        for f in files:
            process(f.read())
    # All files are closed here

Wrapping up

We've covered:

• with statement - Automatic setup and cleanup
• enter and exit - Class-based context managers
• @contextmanager - Generator-based context managers
• Exception handling - Cleanup even on errors
• contextlib - Utilities for common patterns

Key takeaways

1. Use context managers for any setup/cleanup pattern
2. Prefer @contextmanager for simple cases
3. Use class-based for complex state management
4. Always test cleanup happens, even with exceptions
5. Return False from __exit__ unless you want to suppress exceptions

Context managers make your code cleaner and safer. The with statement ensures resources are properly managed, making it harder to introduce resource leaks.