Python Context Managers

In this article, we discuss what context managers are and how you can use them in different ways. There’s a lot of examples we build on to demonstrate different scenarios, and by the time you’re done reading this you will definitely be comfortable using context managers and implementing them in your code!

Let’s just jump right in, shall we?

If you’ve ever opened a file in Python, you probably did something like this:

f = open('foo.txt', 'w')
f.write('Hello, World!')
f.close()

This approach assumes that the we can write to the file without any hiccups, but if an exception does occur during writing, the f.close() statement will never get executed! Also, we have the burden of having to explicitly close the file, and there is a chance we can forgot to do that if our code is complicated. In either case, we can leak a file descriptor in our program if the file never gets closed…oopsie!

To handle any possible exceptions, we can stick our code in a try...except...finally block, ensuring that we close the file in the finally block:

f = open('foo.txt', 'w')
try:
    f.write('Hello, World!')
except Exception as e:
    print(e)
finally:
    f.close()

This is great! Our code handles exceptions, and we are guaranteed that the file will be closed no matter if an exception is thrown.

Let’s take a step back for a moment – resource management (setting up and cleaning up a resource) is a common programming pattern. In our example, cleaning up the resource means closing the file. If we are connecting to a database, we might use this same pattern to ensure that the database connection is closed after we’re done using it. Like many things in programming, patterns can and should be abstracted away. We see this in the next example.

So even though our previous code works, we can trim it down even further using the with statement, like this:

with open('foo.txt', 'w') as f:
    f.write('Hello, World!')

But where did our f.close() statement go? Well, with the magic of the with statement and these things called context managers, our file is being closed but the code that executes the close statement is abtracted away from us.

What are context managers?

Context managers are simply a protocol that an object needs to follow in order for it to be used with the with statement. That “protocol” means that the object needs to define two methods: an __enter__() and an __exit__(). Python docs here.

When we use the with statement, we must call a class or a function that returns a context manager object.

Let’s go back to our file example. Python’s open function returns a context manager, which is then assigned to the variable f. If we do a simple print statement, we can see what our context manager is:

with open('foo.txt') as f:
  print(type(f))  # <class '_io.TextIOWrapper'>

This means that the _io.TextIOWrapper class must have __enter__ and __exit__ methods defined somewhere on its source code. To prove this, let’s inspect the _io.TextIOWrapper class with the dir method to see its attributes and methods:

>>> import _io
>>> dir(_io.TextIOWrapper)
['_CHUNK_SIZE', '__class__', '__del__', '__delattr__', '__dict__', '__dir__', '__doc__', '__enter__', '__eq__', '__exit__', '__format__', '__ge__', '__getattribute__', '__getstate__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__iter__', '__le__', '__lt__', '__ne__', '__new__', '__next__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__sizeof__', '__str__', '__subclasshook__', '_checkClosed', '_checkReadable', '_checkSeekable', '_checkWritable', '_finalizing', 'buffer', 'close', 'closed', 'detach', 'encoding', 'errors', 'fileno', 'flush', 'isatty', 'line_buffering', 'name', 'newlines', 'read', 'readable', 'readline', 'readlines', 'seek', 'seekable', 'tell', 'truncate', 'writable', 'write', 'writelines']

And in fact we do see our two special methods! 🎉 (scroll right)

Fun Fact: If you want to dig into Python’s open function implementation, check it out here in the io module implentation inside the python/cpython repository on GitHub. It is worth noting that the _io module provides the C code that the io module uses internally, so the io module references the _io module.

import _io
import io

print(type(_io.TextIOWrapper))  # <class '_io.TextIOWrapper'>

print(type(io.TextIOWrapper))  # <class '_io.TextIOWrapper'>

In the open method docstring, it says that the method returns a TextIOWrapper. You’ll see that the TextIOWrapper class does not actually define the __enter__ and __exit__ methods on itself but inherits them from the TextIOBase class which inherits them from the IOBase abstract base class.

In fact, to answer our earlier question of where did our f.close() statement go, here is where it is called in the IOBase class.

Source: github.com

More about __enter__ and __exit__

The __enter__ method in a context manager returns the resource (for example, a file object) or returns self, and the __exit__ method does any cleaning up (for example, closing a file, handling exceptions). Of course the class itself can have its __init__ method to set up stuff as well as other class methods.

Some important things to understand here is that when we use the with statement, any code inside the with statement block is executed after the __enter__ method gets called. If an exception is raised within the __init__ or __enter__ methods, the code within the with statement block will not get executed and the __exit__ method will never get called. Once the with statement block is entered though, the __exit__ method is always called and will handle any exceptions.

Let’s take a look at a trivial example:

class DummyContextManager:
    def __init__(self):
        print('__init__')

    def __enter__(self):
        print('__enter__')
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        print('__exit__')

        # Exception description
        print('type, {}'.format(exc_type))
        print('value, {}'.format(exc_value))
        print('traceback, {}'.format(exc_traceback))


with DummyContextManager() as c:
    print('within code block')
print('done')

will give the following output:

__init__
__enter__
within code block
__exit__
type, None
value, None
traceback, None
done

Great! We can clearly see the context flow with these print statements. From the top, the __init__ method is called followed by the __enter__ method. After this, the context enters the with statement block. After the context leaves the with statement block, it returns to the context manager and the __exit__ method is called. If no exceptions occured when the context exited, then the exc_type, exc_value, exc_traceback arguments of the __exit__ method are None. Finally, the context continues to the statement immediately after the with statement.

Let’s try to raise an exception within the __enter__ method to demonstrate that the code within the with statement block and the __exit__ method will not get executed:

class DummyContextManager:
    def __init__(self):
        print('__init__')

    def __enter__(self):
        print('__enter__')

        # this will raise an exception
        quotient = 1/0

        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        print('__exit__')

        # Exception description
        print('type, {}'.format(exc_type))
        print('value, {}'.format(exc_value))
        print('traceback, {}'.format(exc_traceback))


with DummyContextManager() as c:
    print('within code block')
print('done')

And here is the output:

__init__
__enter__
Traceback (most recent call last):
  File "contextmanagers.py", line 19, in <module>
    with DummyContextManager() as c:
  File "contextmanagers.py", line 9, in __enter__
    quotient = 1/0
ZeroDivisionError: division by zero

As you can see, the exception was raised from within the __enter__ method, and our print statement within the with statement block never got called as well as the __exit__ method and also the print statement outside the with statement block.

Now, let’s modify this example in order to let our __exit__ method handle any exceptions thrown. We’ll move our bad division code inside the with statement block:

class DummyContextManager:
    def __init__(self):
        print('__init__')

    def __enter__(self):
        print('__enter__')
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        print('__exit__')

        # Exception description
        print('type, {}'.format(exc_type))
        print('value, {}'.format(exc_value))
        print('traceback, {}'.format(exc_traceback))

        # Custom exception handling
        if isinstance(exc_value, ZeroDivisionError):
            print('Sorry, bad at math')
        else: 
            print('Sorry for the exception!')

        # Surpress exception
        return True


with DummyContextManager() as c:
    print('within code block')

    # this will raise an exception
    quotient = 1/0

    print('still here')
print('done')

Before we see the output, let’s break down what’s going on here. We moved our bad division code inside with with statement block, and we added an additional print statement immediately after within the same block.

Going up to our DummyContextManager class, everything is the same except for the __exit__ method. First, we print the values of exc_type, exc_value and exc_traceback like we’ve done with the previous examples.

Then, we check if the exc_value parameter is an instance of the ZeroDivisionError exception class. Note that the exc_type parameter is the exception class itself (whatever it may be if an exception occurs), and exc_value is the instance of the exception class. So basically, we’re just printing out some custom message based on what exception occurs. Finally, we return True to surpress the exception and allow our program to continue.

Now, let’s see the output of this:

__init__
__enter__
within code block
__exit__
type, <class ZeroDivisionError>
value, division by zero
traceback, <traceback object at 0x101d59a08>
Sorry, bad at math
done

Interesting! This output should have been expected. Our exception has been caught, and due to our custom exception handling, our apologetic print statement is printed out.

Also notice that our still here print statement never got printed, because the context left the with statement block as soon as the exception occured.

Using contextlib.contextmanager

Defining classes with __enter__ and __exit__ methods is not the only way of writing context managers.

We can use the contextlib.contextmanager decorator to define a generator based function that will work with the with statement. To do this, we create a function the @contextmanager decorator and call yield once. Calling yield means that the code within the decorated function suspends while it “yields” the resource to the user and is resumed after the yield occurs. You can think of it like everything before the yield is the __enter__ method in class-based context managers, and everything after the yield is like the __exit__ method.

Let’s recreate our dummy context manager example using the @contextmanager decorator:

from contextlib import contextmanager

@contextmanager
def dummy_context_manager():
    print('Enter')
    yield
    print('Exit')


with dummy_context_manager() as d:
    print('within code block')

and we get the expected output:

Enter
within code block
Exit

We can also use try...except...finally blocks to handle exceptions. According to the Python docs, any unhandled exceptions in the with statement block are reraised when the generator resumes after the yield statement. So we can use try...except...finally to handle any exceptions that come up and handle any clean up.

Take this example:

from contextlib import contextmanager

@contextmanager
def dummy_context_manager():
    try:
        print('Enter')
        yield
    except Exception as e:
        print('Some exception occured!')
    finally:
        print('Exit')


with dummy_context_manager() as d:
    print('within code block')

    # this will raise an exception
    quotient = 1/0
    
    print('still here')
print('done')

which gives the following output:

Enter
within code block
Some exception occured!
Exit
done

Our exception was nicely handled by the try...except...finally block. From the top, the generator yielded, then the code in the with statement block began executing. The exception from our bad math occured and was reraised inside the generator in the except block. This is where the generator resumes. After the except block handles the exception, the finally block is executed. Then, the execution continues with the statement following the with statement because our print statement after the bad math code within the with statement never gets called.

For another example, lets implement a file handler like the open function using the @contextmanager decorator:

from contextlib import contextmanager

@contextmanager
def file_manager(file_path, mode):
    try:
        open_file = open(file_path, mode)
        yield open_file
    except Exception as e:
        print('Some exception occured!')
    finally:
        open_file.close()


with file_manager('foo.txt', 'w') as f:
    f.write('Hello, World!')

Context Manager Example - Timer

For our final example, we will implement a context manager that measures the execution of a code block with both a class and the decorator method. For this, we will use the time.time function.

First up, the class-based context manager:

import time

class Timer:
    def __init__(self):
        pass

    def __enter__(self):
        print('Start timer...')
        self.start = time.time()
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        self.stop = time.time()
        print('Stop timer...')
        total = self.stop - self.start
        print('Total time: {:.3f} seconds'.format(total))


with Timer() as t:
    time.sleep(5)

and we have our output:

Start timer...
Stop timer...
Total time: 5.003 seconds

Next up, the generator context manager:

from contextlib import contextmanager
import time

@contextmanager
def timer():
    print('Start timer...')
    start = time.time()
    yield
    stop = time.time()
    print('Stop timer...')
    total = stop - start
    print('Total time: {:.3f} seconds'.format(total))


with timer() as t:
    time.sleep(5)

with the following output:

Start timer...
Stop timer...
Total time: 5.001 seconds

Both approaches behave the same. Ultimately, its up to you depending on which approach you’re more comfortable with!

Context managers can be really handy when it comes to manipulating the control flow of a particular resource in your code. Even though the open function context manager is probably the most popular example, other useful context managers have been added to the Standard Library for example, threading.Lock objects. You can always come up with your own context managers too, when it makes sense to use them.

By now, I hope these examples have illustrated what context managers are and just how maintainable they can make your code! Enjoy messing around with some of these examples and happy coding! 😄