Python: are your context managers correctly written?

Context managers are a powerful tool when used properly.

What are context managers?

Even if you have never heard of context managers, you have probably seen them without knowing it. The most typical use case is when they are used to take a resource at the beginning of a block and make sure it is released at the end of it. For example, the open method can (and should) be used as a context manager:

with open(filename) as file:
    ...

Written this way, regardless of what happens within the with block, the file will eventually be closed. It is equivalent to writing:

file = open(filename)
try:
    ...
finally:
    file.close()

Context managers can also be used for many more things. For example, you can also write a custom context manager to easily time a piece of code and send it to your monitoring system:

with monitor_time():
    my_critical_method()

I have been writing and reviewing python code in back-end applications for several years now, and there is one critical mistake that developers often make. As context managers sometimes handle very critical resources such as thread locks or file descriptors, I feel it is important to raise awareness about this pitfall as it can be a source of important downtime or bugs in your applications if left unchecked.

Writing a context manager in Python

Python offers two standard ways to write a custom context manager: a class-based approach and a generator-based approach. Let’s dig into these two approaches by writing a simple timing context manager.

The class-based approach

This first approach consists of creating a class with two methods:

• the __enter__ method that will be called when entering the with block,
• and the __exit__ that will be called when exiting the with block. This method takes 3 arguments that are provided with information about the exception raised in the with block, but we are not interested in them here. All you need to know is they are often called type , value and traceback in this order.

The generator-based approach

This second approach is more straightforward. It consists of writing a method with one yield that marks when the content of the with block will be executed. This method may execute some code before and after this yield to perform its logic. In our case, it would look something like this:

This generator-based approach works fine most of the time, but it has one pitfall you must be careful about. We will see what it is and how to write this generator_based_monitor_time context manager correctly when using the generator-based method.

The pitfall of the generator-based method

Testing our context managers

Let’s now try our brand new context managers! First, with a body that performs a very complex computation that we want to monitor:

This prints:

Entering class-based context manager
sleep: 1.00 seconds elapsed
Left class-based context manager
Entering generator-based context manager
sleep: 1.00 seconds elapsed
Left generator-based context manager

As expected, we see that both context managers print the elapsed time after the sleep function finished. But now what happens if the with body raises an exception?

This prints:

Entering class-based context manager
exception: 0.00 seconds elapsed
Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
RuntimeError
Entering generator-based context manager
Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
RuntimeError

Oh! We see that with our second context manager the timing print was not executed!

Why is this happening? And how can we fix this?

Now in our monitoring case, this print not being shown is probably not a very important issue. But context managers are often used in much more critical cases to make sure resources are properly released. For example, if a badly written context manager was used to take a lock, then if an exception occurs and the context manager does not exit properly, the lock will never be released. This in turn may be much more problematic and cause unexpected deadlocks.

So now, how can we fix our second context manager? One way to understand what happens in this case is to replace the yield of our context manager with the content of the with block that uses it. Even if that’s not exactly what happens under the hood technically, this visualization is still very helpful to understands what happens. So in our case this gives us:

start_time = time.perf_counter()
raise RuntimeError
elapsed_time = time.perf_counter() - start_time
print(f"{name}: {elapsed_time:.2f} seconds elapsed")

Here we clearly see that the second part of the context manager will never be executed, as the error is raised and the function exits before executing its cleanup code.

But now looking at the code this way, we can see how we should fix our context manager: we can add a try... finally... block around the yield so that we are immune to exceptions raised in the with block.

@contextmanager
def generator_based_monitor_time_2(name):
    start_time = time.perf_counter()
    try:
        yield
    finally:
        elapsed_time = time.perf_counter() - start_time
        print(f"{name}: {elapsed_time:.2f} seconds elapsed")

Now let’s check our new context manager exits properly:

print("Entering 2nd generator-based context manager")
with generator_based_monitor_time_2("exception"):
    raise RuntimeError
print("Left generator-based context manager")

This gives us:

Entering 2nd generator-based context manager
exception: 0.00 seconds elapsed
Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
RuntimeError

Perfect! Now our context manager does close properly in all cases, even if an exception is raised!

Conclusion

In this article, we learned what context managers are and the two syntaxes we can use to write our own. Of those two syntaxes, the generator-based one is often preferred as it is easier to write, but there is one critical pitfall you must avoid: you should always be careful about wrapping its yield in a try... finally... block to make sure it always releases its resources properly. Getting into the habit of always doing so, even in non-critical cases like the one used as an example in this article, will help you remember it when the context manager you write will be handling much more critical resources. For reference, if you need help remembering the syntax, the python documentation provides a template example that you can adapt to fill all your needs:

@contextmanager
def managed_resource(*args, **kwds):
    # Code to acquire resource, e.g.:
    resource = acquire_resource(*args, **kwds)
    try:
        yield resource
    finally:
        # Code to release resource, e.g.:
        release_resource(resource)

About us

Our mission at Mindsay is to help companies provide simple and efficient interactions to all of their customers at any time.
Check us out!

We’re also hiring