Context Managers

python
code
today I learned
Author

Mat Miller

Published

December 13, 2022

Open in Google Colab
This is a quick ‘today I learned’ (TIL) note on Python Context managers. Python context managers are used to wrap arbitrary code with entry (setup) and exit (cleanup) functions. One common places you’ll see them used is when reading data from a file.

# Open file and read contents.
with open('test.txt','r') as f:
    output = f.readlines()
print(output)
['This file is called test.txt.\n', "This is what's on the second line."]

If we try and read from the file f, defined above, we will get an I/O exception because the file as already been closed.

try:
    f.readlines()
except Exception as e:
    print(e)
I/O operation on closed file.

Here is the equivalent long hand way to read the data from the file:

f = open('test.txt')
output = f.readlines()
f.close()
print(output)
['This file is called test.txt.\n', "This is what's on the second line."]

As you can see the syntax is more verbose, it would be easier to forget to close the file, and it’s much less clear to see at a glance when we’re operating on the file. This example is relatively trivial as we’re just reading all the lines of the text file into a list but you can probably imagine this could be a lot more complex if you were doing something more complicated like training a neural net.

Now let’s write our own class that uses a conext manager to cement how they can be implemented.

class MyContextManagerClass:
    def __enter__(self):
        print("Entering the context...")
        return "My enter message."
    def __exit__(self, exception_type, exception_value, exception_traceback):
        print("Leaving the context...")
        print(exception_type, exception_value, exception_traceback, sep="\n")
with MyContextManagerClass() as h:
    print('hi', h)
Entering the context...
hi My enter message.
Leaving the context...
None
None
None

As you can see the enter message was printed, the __enter__ return value was passed and then the exit message was printed. Now let’s see what happens if there is an error while within our context.

with MyContextManagerClass() as h:
    print(h)
    print(1/0)
Entering the context...
My enter message.
Leaving the context...
<class 'ZeroDivisionError'>
division by zero
<traceback object>
ZeroDivisionError: division by zero

As you can see an error was thrown but the __exit__ function was run anyways.

There are many other ways you can implement and use context managers which you can read about here: Python Conext Managers. Hopefully I’ve given you a taste of what’s possible and given you a basic understanding of they they’re useful.

Here are a few more examples for your reference:

Example 1: Using the contextmanager decorator

from contextlib import contextmanager

@contextmanager
def closing(thing):
    try:
        print('Starting')
        yield thing
    finally:
        print('Finishing:',thing)
with closing('a'):
    print('hi')
Starting
hi
Finishing: a

Example 2: Using ContextDecorator

from contextlib import ContextDecorator

class mycontext(ContextDecorator):
    def __enter__(self):
        print('Starting')
        return self

    def __exit__(self, *exc):
        print('Finishing')
        return False
@mycontext()
def my_function():
    print('The bit in the middle')
my_function()
Starting
The bit in the middle
Finishing