Working With Context Managers in Python

Assuming you are someone who has been writing python for at least a couple of weeks, you’ve probably come across lines of code beginning with a with keyword. Perhaps you were even encouraged to use them in certain areas of your code but didn’t know why. These statements give access to what is called a Context Manager. In this article, I dive into what they are, why you should use them and how you can implement them in your python code.

About Context Mangers…

Let’s start with what a context manager looks like when it’s implemented. Traditionally, if you wanted to, say, print out the first line of a file you’d probably do something like this:

file = open('my_file.txt', 'r')
try:
    print(file.readline())
finally:
    file.close()

This approach is mostly fine, but it means that each time we choose to do it this way, we always need to remember to implement some functionality for when the code exits. In this case, we have to pay attention to close the opened file using file.close() which also comes with an additional baggage of needing to use the try ... except ... finally commands to guard against errors and ensure our file always closes.

Context managers effectively allow us to say: “Before running the block of code C, do X. Also, do Y after, whether C was successful or not.”. When we implement the context manager in the open() function for accessing a file, it implicitly closes the file at the exit point of the code block so we don’t always have to. Therefore, we can write an equivalent of the above code in two lines:

with open('my_file.txt', 'r') as file:
    print(file.readline())

Why use them?

Avoids resource leaks

For each opened file or resource, there needs to be an equivalent close operation to clean up. Failure to do this results in the memory leaks where allocated memory isn’t freed even though it is no longer needed.

When your parents leave your room without closing the door

I imagine this ☝️ is how your RAM feels when opened files don’t get closed.

Memory leaks can be especially bad when the faulty block of code is being run continuously or in environments that do not terminate. Leaks could accumulate and it could end up significantly slowing down the performance of the system or in the worst case, causing a complete system failure.

Succinct and easy to read

The with statement as demonstrated in the file reading example above, lets us write significantly shorter code in an intuitive pythonic way. It takes away the necessary additional operations (like closing of a file) and performs them under the hood so that at point of implementation, you only need to worry about the main operations you want to carry out within that context.

Handles Exceptions

When exceptions occur within a context scope, context managers can handle them so you don’t always have to. In the file reading example above, we no longer needed to anticipate exceptions with a try ... finally clause when we used the with statement. This is already handled in the underlying implementation of the in-built open() function.

Common applications of context managers are in acquiring and releasing locks in thread synchronisation, managing database connections and accessing files (as seen above). However, suppose you have custom operations that you want to execute as a pair with some code between them, you could implement your own custom context manager.

Implementing a custom Context Manager

Your user-defined classes can also have runtime contexts which are executed using the with statement. For this to be possible, the class must implement at least the __enter__ and __exit__ methods.

• __enter__(): This is the entry point of the runtime context. Here you can perform operations that you want to happen at the start of the runtime context. You can also optionally return a value which is bound to the identifier in the as clause of the with statement. The __enter__() method is called after the __init__() method which still serves its purpose as the constructor called when an object is created.
• __exit__(type, value, traceback): This is the method called when the runtime context exits. Hence, this is the place to perform operations that you want to be done at the end of the context. To help with this, it provides 3 arguments for handling exceptions that may occur while in the context body; the exception type, the value and traceback. As you may imagine, these values are None if no exceptions occur. If a True is returned in this method and an exception occurs in the with statement body at runtime, the error is silenced and the code execution is resumed from after the with block. A False return value propagates the exception as usual.

Here’s a simple example of a custom context manager which calculates the time spent to execute the block of code within the context.

import time

class TimedOperation:
    def __init__(self, operation):
        self.operation = operation
    def __enter__(self):
        self.start_time = time.time()
        print('Starting Operation "{}"'.format(self.operation))
    def __exit__(self, type, value, traceback):
        time_elapsed = time.time() - self.start_time
        print('Operation "{}" ran for {} seconds'.format(self.operation, time_elapsed))

Implementing this at runtime to calculate the time spent to sum up the first 1 million numbers, we have:

with TimedOperation('Adding numbers'):
    x = 0
    for i in range(1000000):
        x += i

Running this piece of code gives the following output:

Starting Operation "Adding numbers"
Operation "Adding numbers" ran for 0.09711217880249023 seconds

Context managers as seen above can be very useful. Embracing the idea and knowing when and how to use them can save you a lot of time and resources and frankly, make you write more ‘pythonic’ code. Hopefully, you’ll be coming out of the context of this article on a successful run and a read that was free from exceptions.