Writing your own contextmanager using generator syntax

suggest change

It is also possible to write a context manager using generator syntax thanks to the contextlib.contextmanager decorator:

import contextlib

@contextlib.contextmanager
def context_manager(num):
    print('Enter')
    yield num + 1
    print('Exit')

with context_manager(2) as cm:
    # the following instructions are run when the 'yield' point of the context
    # manager is reached.
    # 'cm' will have the value that was yielded
    print('Right in the middle with cm = {}'.format(cm))

produces:

Enter
Right in the middle with cm = 3
Exit

The decorator simplifies the task of writing a context manager by converting a generator into one. Everything before the yield expression becomes the __enter__ method, the value yielded becomes the value returned by the generator (which can be bound to a variable in the with statement), and everything after the yield expression becomes the __exit__ method.

If an exception needs to be handled by the context manager, a try..except..finally-block can be written in the generator and any exception raised in the with-block will be handled by this exception block.

@contextlib.contextmanager
def error_handling_context_manager(num):
    print("Enter")
    try:
        yield num + 1
    except ZeroDivisionError:
        print("Caught error")
    finally:
        print("Cleaning up")
    print("Exit")

with error_handling_context_manager(-1) as cm:
    print("Dividing by cm = {}".format(cm))
    print(2 / cm)

This produces:

Enter
Dividing by cm = 0
Caught error
Cleaning up
Exit

Found a mistake? Have a question or improvement idea? Let me know.

Table Of Contents