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:
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
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:
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:
This means that the
_io.TextIOWrapper class must have
__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:
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
open method docstring, it says that the method returns a TextIOWrapper. You’ll see that the
TextIOWrapper class does not actually define the
__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
__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
__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:
will give the following output:
Great! We can clearly see the context flow with these
__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_traceback arguments of the
__exit__ method are
None. Finally, the context continues to the statement immediately after the
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:
And here is the output:
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:
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_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:
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.
Defining classes with
__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
Let’s recreate our dummy context manager example using the
and we get the expected output:
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:
which gives the following output:
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
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
First up, the class-based context manager:
and we have our output:
Next up, the generator context manager:
with the following output:
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! 😄