Skip to main content Brad's PyNotes

Context Managers: Mastering Python's Resource Management

TL;DR

Context managers handle setup and cleanup automatically using the with statement, implementing __enter__ and __exit__ methods to guarantee resource cleanup even when exceptions occur.

Interesting!

Context managers can suppress exceptions by returning True from __exit__, allowing you to selectively handle errors without cluttering your code with try-except blocks.

The Context Manager Protocol

Context managers implement two special methods that Python calls automatically:

python code snippet start

class DatabaseConnection:
    def __init__(self, db_name):
        self.db_name = db_name
        self.connection = None

    def __enter__(self):
        print(f"Opening connection to {self.db_name}")
        self.connection = connect_to_database(self.db_name)
        return self.connection  # This becomes the 'as' variable

    def __exit__(self, exc_type, exc_value, traceback):
        print(f"Closing connection to {self.db_name}")
        if self.connection:
            self.connection.close()
        return False  # Don't suppress exceptions

# Usage
with DatabaseConnection("mydb") as conn:
    conn.execute("SELECT * FROM users")
# Connection automatically closed

python code snippet end

The __exit__ method receives three arguments about any exception that occurred:

  • exc_type: The exception class (or None)
  • exc_value: The exception instance (or None)
  • traceback: The traceback object (or None)

Creating Context Managers with Decorators

The contextlib.contextmanager decorator simplifies context manager creation:

python code snippet start

from contextlib import contextmanager
import time

@contextmanager
def timer(name):
    start = time.time()
    print(f"Starting {name}")
    try:
        yield  # Code block executes here
    finally:
        elapsed = time.time() - start
        print(f"{name} took {elapsed:.2f}s")

# Usage
with timer("data processing"):
    process_large_dataset()

python code snippet end

The code before yield acts as __enter__, the code after acts as __exit__, and the finally block ensures cleanup happens.

Exception Handling in Context Managers

Context managers can catch and suppress exceptions:

python code snippet start

class IgnoreErrors:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is ValueError:
            print(f"Caught and suppressed: {exc_value}")
            return True  # Suppress the exception
        return False  # Propagate other exceptions

with IgnoreErrors():
    int("not a number")  # ValueError caught and suppressed
    print("This runs because exception was suppressed")

with IgnoreErrors():
    raise TypeError("Different error")  # This propagates

python code snippet end

The contextlib Toolkit

suppress() - Ignore Specific Exceptions

python code snippet start

from contextlib import suppress

# Instead of try-except
with suppress(FileNotFoundError):
    os.remove('temp_file.txt')

# Suppress multiple exception types
with suppress(KeyError, AttributeError):
    value = config[key].attribute

python code snippet end

redirect_stdout() and redirect_stderr()

python code snippet start

from contextlib import redirect_stdout, redirect_stderr
import io

# Capture printed output
buffer = io.StringIO()
with redirect_stdout(buffer):
    print("This goes to the buffer")
    help(len)

output = buffer.getvalue()

python code snippet end

ExitStack - Dynamic Context Managers

python code snippet start

from contextlib import ExitStack

# Open multiple files dynamically
filenames = ['file1.txt', 'file2.txt', 'file3.txt']

with ExitStack() as stack:
    files = [stack.enter_context(open(fname)) for fname in filenames]
    # All files automatically closed at the end
    for f in files:
        print(f.read())

# Conditional cleanup
with ExitStack() as stack:
    stack.callback(cleanup_temp_files)
    result = risky_operation()
    if result.success:
        stack.pop_all()  # Don't cleanup on success

python code snippet end

nullcontext() - Optional Context Management

python code snippet start

from contextlib import nullcontext

def process_data(data, use_lock=True):
    lock = threading.Lock() if use_lock else nullcontext()

    with lock:
        # Works whether lock is real or no-op
        modify_shared_data(data)

python code snippet end

closing() - Ensure Close Called

python code snippet start

from contextlib import closing
from urllib.request import urlopen

with closing(urlopen('https://www.python.org')) as page:
    for line in page:
        process(line)
# urlopen result automatically closed

python code snippet end

Asynchronous Context Managers

Async context managers use __aenter__ and __aexit__ with async with:

python code snippet start

class AsyncDatabaseConnection:
    async def __aenter__(self):
        self.conn = await connect_async()
        return self.conn

    async def __aexit__(self, exc_type, exc_value, traceback):
        await self.conn.close()
        return False

# Usage
async def fetch_data():
    async with AsyncDatabaseConnection() as conn:
        return await conn.fetch("SELECT * FROM users")

python code snippet end

Using the decorator approach:

python code snippet start

from contextlib import asynccontextmanager

@asynccontextmanager
async def async_timer(name):
    start = time.time()
    try:
        yield
    finally:
        print(f"{name}: {time.time() - start:.2f}s")

async def main():
    async with async_timer("API call"):
        await fetch_from_api()

python code snippet end

Practical Patterns

Temporary State Changes

python code snippet start

@contextmanager
def temporary_attribute(obj, attr, value):
    old_value = getattr(obj, attr)
    setattr(obj, attr, value)
    try:
        yield
    finally:
        setattr(obj, attr, old_value)

# Usage
config = Config(debug=False)
with temporary_attribute(config, 'debug', True):
    # Debug mode temporarily enabled
    run_tests()
# Debug mode restored

python code snippet end

Transaction Management

python code snippet start

@contextmanager
def transaction(connection):
    try:
        yield connection
        connection.commit()
    except Exception:
        connection.rollback()
        raise

with transaction(db_conn) as conn:
    conn.execute("UPDATE accounts SET balance = balance - 100")
    conn.execute("UPDATE accounts SET balance = balance + 100")
# Automatically commits on success, rolls back on error

python code snippet end

Multiple Context Managers

python code snippet start

# Multiple contexts in one statement
with open('input.txt') as infile, open('output.txt', 'w') as outfile:
    outfile.write(infile.read().upper())

# Nested contexts for clarity
with ExitStack() as stack:
    input_files = [stack.enter_context(open(f)) for f in inputs]
    output_file = stack.enter_context(open('combined.txt', 'w'))

    for f in input_files:
        output_file.write(f.read())

python code snippet end

Best Practices

  1. Always return from __exit__: Explicitly return False to avoid accidentally suppressing exceptions
  2. Use finally in decorators: Ensure cleanup code runs even if exceptions occur
  3. Keep context managers focused: Each should manage one resource or concern
  4. Prefer contextlib for simple cases: Use @contextmanager instead of writing classes for simple cleanup
  5. Make context managers reusable: Design them to work multiple times unless single-use is required

Context managers are the Pythonic way to handle resources and ensure cleanup. They integrate seamlessly with the with statement and work alongside exception handling . Use them with database connections , thread locks , and file operations . For async code, combine with asyncio patterns.

Reference: Python Data Model - With Statement Context Managers and contextlib - Utilities for with-statement contexts