Building a context manager from scratch.

Let's consider this code example one more time.

with open_db_connection() as connection:
    run_analysis(connection)

This code contains a with-block. And once the with block is entered, the connection gets created and stored in the connection variable. Then, the body of the with block gets executed, after which the connection will get closed by the context manager automatically, no matter whether run_analysis threw an exception or not.

This with statement relies in the back on two special methods: object.__enter__(self) and object.__exit__(self, exc_type, exc_value, traceback)

Any object that implements these two methods can be used in the header of a with statement. Let's look at an example:

class PrintingContextManager:
    def __enter__(self):
        print('entering the context manager')

    def __exit__(self, exc_type, exc_value, traceback):
        print('exiting the context manager')


with PrintingContextManager():
    print('inside the context manager')

# This prints:
# entering the context manager
# inside the context manager
# exiting the context manager

You can see that we enter the block, run some code and exit it again. You can even see the enter and exit print statements when we introduce an error.

class PrintingContextManager:
    def __enter__(self):
        print('entering the context manager')

    def __exit__(self, exc_type, exc_value, traceback):
        print('exiting the context manager')


with PrintingContextManager():
    raise ValueError('oh no!')

We can even use the exc_type, exc_value and traceback variables in the __exit__ method to inspect useful information about the exception. This is all information that might help you shut down elegantly.

class PrintingContextManager:
    def __enter__(self):
        print('entering the context manager')

    def __exit__(self, exc_type, exc_value, traceback):
        print('exiting the context manager')
        print(exc_type)
        print(exc_value)
        print(traceback)


with PrintingContextManager():
    raise ValueError("oh no!")

Finally, you're also able to have the __enter__ method return a value. This allows you to use the with <manager> as <variable>: syntax.

class PrintingContextManager:
    def __enter__(self):
        print('entering the context manager')
        return "i am the returned value"

    def __exit__(self, exc_type, exc_value, traceback):
        print('exiting the context manager')


with PrintingContextManager() as var:
    print('inside the context manager')
    print(var)