Context Managers

Have you ever opened a file using a with statement and marvelled at how it automatically closes itself when you’re done?

The magic behind the with statement is the Context Manager Protocol. A context manager is an object that defines a runtime context by handling the entry into and exit from a block of code.

Here is a visual representation of the context management lifecycle:

Let’s trace how execution flows from __enter__ into the block, and finally to __exit__. Change value to see how exceptions trigger early exit!

class SuppressErrors:
    def __enter__(self):
        print("Entering context")
        return self
        
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            print(f"Exception suppressed: {exc_type.__name__}")
            return True # Suppress the exception!
        print("Exiting normally")

value = 0

with SuppressErrors():
    print(f"Working with value: {value}")
    if value == 0:
        raise ValueError("Cannot be zero!")
    print("Work complete")
print("Execution continues safely")

Building Class-Based Context Managers

To build a custom context manager using a class, you just need to implement __enter__ and __exit__.

class Timer:
    def __enter__(self):
        print('Timer started')
        return self
    def __exit__(self, exc_type, exc_val, exc_tb):
        print('Timer stopped')

with Timer() as t:
    print('Doing work')

Timer started
Doing work
Timer stopped

Suppressing Exceptions

The __exit__ method takes three arguments related to an exception: the exception type, value, and traceback. If the context was exited normally, all three are None.

If an exception occurs and __exit__ returns a truthy value, Python suppresses the exception!

class SuppressErrors:
    def __enter__(self):
        return self
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            print(f'Suppressed {exc_type.__name__}')
            return True

with SuppressErrors():
    raise ValueError('Oops')
print('Execution continues')

Suppressed ValueError
Execution continues

The @contextmanager Decorator

Writing classes just to get __enter__ and __exit__ can be boilerplate-heavy. Python’s contextlib module provides the @contextmanager decorator, which turns a generator function into a full context manager.

Everything before the yield statement acts as __enter__. The yield itself passes the bound value to the as clause. Everything after the yield acts as __exit__.

from contextlib import contextmanager

@contextmanager
def tag(name):
    print(f'<{name}>')
    yield
    print(f'</{name}>')

with tag('h1'):
    print('Hello World')

<h1>
Hello World
</h1>

Exception Propagation in Generators

When using @contextmanager, any unhandled exception in the with block is literally re-raised inside your generator at the line where yield occurred. To handle cleanup safely, you must wrap your yield in a try/finally block!

from contextlib import contextmanager

@contextmanager
def safe_open():
    print('Opening')
    try:
        yield
    except ValueError:
        print('Caught ValueError')
    finally:
        print('Closing')

with safe_open():
    print('Working')
    raise ValueError('Bad value')

Opening
Working
Caught ValueError
Closing

Managing Dynamic Contexts: ExitStack

Sometimes you don’t know ahead of time how many context managers you need to open (e.g., dynamically opening multiple files). The ExitStack is designed exactly for this: it allows you to dynamically register context managers, and it ensures all of them are cleaned up properly when the stack is exited.

from contextlib import ExitStack

with ExitStack() as stack:
    print('Managing multiple contexts dynamically')
    # stack.enter_context(open('file.txt')) would go here
    print('Done')

Managing multiple contexts dynamically
Done

The Standard Library Example

The most prominent example of a context manager is file handling. File objects natively implement this protocol.

import tempfile
import os

with tempfile.NamedTemporaryFile(delete=False) as f:
    f.write(b'Hello')
    name = f.name

with open(name, 'r') as f:
    print(f.read())

os.remove(name)

Hello