Up Your Python Coding Skills: Context Managers

Some of the more common operations we’ll ever perform in our day-to-day activities as Software Engineers are managing external resources, like files or database connections, just to name a few. These resources need to be managed carefully, in the sense that they should be released after usage, or we’ll be risking resource leakage, leading to slow, unresponsive systems that may even crash on us. To be more specific, take the example of a file operation: we’d write a code that does a wonderful job at opening a file, reading from it, or writing to it, but that omits the closing of said file. Next time this code would run, we’d have two files open. See the problem?

In order to achieve this proper resource management, we’d need to ensure that every resource benefits from an initial, setup phase — this would be the opening of the file — and a final, teardown phase — the successful closing of the file.

As such, let’s consider the following three lines of code:

# resource initialization
my_file = open("test.txt", "w")

# resource processing
my_file.write("This is a test")

# resource teardown
my_file.close()

In effect, we’ve ensured we got the indispensable setup-teardown duo that ensures the proper releasing of the resource as soon as it’s not needed anymore. Or did we?

There’s nothing to suggest the file will be successfully closed if an exception occurs during the writing process. Should that happen, we’d obviously never have the close() method called on the file and we’d be back at square one: resource leakage.

One way to deal with this problem is to use the try...except...finally approach:

# resource initialization
my_file = open("test.txt", "w")

try:
    # resource processing
    my_file.write("This is a test")
except Exception as exc:
    print(f"An exception occurred: {exc}")
finally:
    # resource teardown
    my_file.close()

The finally keyword in Python ensures that the lines in that block will execute regardless of whether we had an exception thrown in the try block — and caught in the except block or not. So any statement present under finally will execute, exception or not.

This seems to have solved the problem quite nicely. And we’d be right to assume so. But perhaps there’s an even better way to achieve this? I mean, it still relies a bit too much on our memory/experience to not forget/omit to use this construct.

This is where Python comes to our aid, via these constructs called context managers.

A context manager, in Python, is an object (like most everything in Python, for that matter) that sets up a context for code to run in. Neatly hidden behind the elegant, yet still kind of overlooked with statement, the created context provides that much-needed setup-teardown duo we were mentioning just now that would ensure proper initialization and release of the external resource, preventing memory leaks or worse. Here’s the equivalent of our try-except-finally example from earlier, this time written using the with statement:

# resource initialization - context is being created
with open("test.txt", "w") as my_file:
    # resource processing
    my_file.write("This is a test")
    
# resource teardown happens as soon as we exit the context
# context is cleaned automatically
# but we can also access the context after we're done with it
print(my_file.closed)

Output:
True

That is a very elegant and neat solution to the proper handling of external resources. What’s more interesting is that we’re not by any means confined to using the built-in context managers Python has to offer. We can build our very own context managers and I think everyone should be at least a tiny bit interested in this, as it sheds so much light on what’s really happening behind the scenes.

If we want a class that we’re building to also have context manager capabilities, we need only implement two methods:

__enter__()
__exit__()

The first method, __enter__() is responsible for setting up and initialization of the context. It can (but it is not a mandatory requirement) return some object.

The second method, __exit__() is where the cleanup takes place.

Typically, what happens is this:

with context_manager() as context:
    context.some_method()

1. We write the with statement. Here’s where the context manager’s being created: with context_manager() as context; what happens is __init__() of the context manager is called;
2. Immediately after the creation of the context manager, __enter__() is called. If __enter__() also returns an object, that object is being assigned to my_file;
3. Processing occurs: context.some_method();
4. Once we’re done processing the code under the with block, Python exits the context; what really happens is __exit__() is called; it is very important to note that __exit__() is being called even if some exception has occurred in the processing block of code.

The general structure of a class implementing a context manager looks kinda like this:

class MyContextManager():
    def __init__(self):
        print("Called init method")

    def __enter__(self):
        print("Called enter method")

    def __exit__(self, exc_type, exc_value, exc_traceback):
        print("Called exit method")

with MyContextManager() as manager:
    print("Inside the with block")

Output:
Called init method
Called enter method
Inside the with block
Called exit method

You can see the four steps we just enumerated earlier and, more importantly, notice the order in which these statements are being executed.

A very nice to know aspect is this: we can nest context managers. Say we have 2 separate context manager classes, called Context_1 and Context_2. If, for some reason, we’d need to use them in a nested fashion, all we’d need to do at this point would be this:

with Context_1() as manager_1:
    with Context_2() as manager_2:
        process_stuff()

Now, maybe it’s time to address some natural questions you may have at this point: if the syntax for the with statement is typically of this form: with context_manager() as context, then why does our previous example where we had with open("test.txt", "w") as my_file work? I mean, open() isn’t a class, right?

Yes, you’d be right to ask this question. I’ve asked it myself too, back in the day. Turns out, open() isn’t a class, it’s still a function, but its result is a file object that implements both the __enter__() and the __exit__() methods. Effectively, in this case, it’s the result of the open() function that is the actual context manager. This is why with open("test.txt", "w") as my_file works.

Needless to say, this article doesn’t even begin to do Python’s context managers any justice. Nor does it aims to do so. It’s merely an induction-type article, where the most basic features and traits are being presented and, as such, it’s solely meant to familiarize the reader with a couple of things that make Python one of the best programming languages to work with.

In the end, I’ll leave you with yet another side to this context manager construct: the contextlib module, which provides us with very useful utilities for working with context managers and, by extension, with the with statement. A great starting point is, obviously, represented by the official docs.

Happy coding and stay safe, guys! Until next time!

Deck is a software engineer, mentor, writer, and sometimes even a teacher. With 12+ years of experience in software engineering, he is now a real advocate of the Python programming language while his passion is helping people sharpen their Python — and programming in general — skills. You can reach Deck on Linkedin, Facebook, Twitter, and Discord: Deck451#6188, as well as follow his writing here on Medium.

More content at PlainEnglish.io. Sign up for our free weekly newsletter. Follow us on Twitter and LinkedIn. Check out our Community Discord and join our Talent Collective.