Embracing python context managers

Table of contents

Most languages have the tried and tested try ... catch/except ... finally concept of dealing with resources that require cleanup, such as files, network sockets or database connections. While python offers that feature, it also includes a keyword that can be seen as a direct competitor to it: the with statement.

Why the with syntax?

To understand why you might prefer the with style over try/finally, let's look at some basic file operations:

f = None
   f = open("test.txt")
   if f is not None:

This is fairly standard syntax you should recognize even from other languages. One little caveat is that we need to define variable f before the call to open(), because if that throws an exception then f won't be defined, thus we cannot check it in the finally block.

Now let's write the same code using the with statement:

with open("test.txt") as f:

That's a lot shorter! You can even use multiple context managers with a single statement:

with open("test.txt") as f1, open("test2.txt") as f2:

How do context managers work?

Now that we have seen how many lines of code the with statement saves us, let's find out how the magic happens. At it's core, it's really quite simple: The with statement takes a context manager. A context manager is just an object that contains at least these two methods: __enter__() and __exit__().

Here is an example to help you see it in action:

class Car(object):
   def __init__(self, name):
      self.name = name
   def __enter__(self):
      print(f"__enter__ was called for {self.name}")
      return self
   def __exit__(self, type, value, traceback):
      print(f"__exit__ was called for {self.name}")

with Car("toyota") as my_car:
   print(f"Driving my {my_car.name}")
   raise Exception("Traffic stop")

This prints the following:

__enter__ was called for toyota
Driving my toyota
__exit__ was called for toyota
Traceback (most recent call last):
  File "/tmp/py/main.py", line 12, in <module>
    raise Exception("Traffic stop")
Exception: Traffic stop

As you can see, the __enter__() method is called before any code from our with block runs, and the __exit__() method is invoked implicitly before the exception is thrown, just like a finally block would have done.

It is important to note that __enter__() doesn't need to return anything, and we don't have to use the return value. Even if we simply remove the as my_car part from the with statement, the it would still keep track of our context manager class and call __enter__() and __exit__() at the correct moment. The arguments type, value and traceback for the __exit__() method refer to the exception that was raised in the with block (if any).

Saving boilerplate with contextlib

Now that we have seen how to make your own context manager, let's reduce the amount of boilerplate code required to create them. To do this, python's contextlib module provides the convenience decorator contextlib.contextmanager. By using this decorator, we can turn a simple generator function into a complete context manager. Rewriting our code above with this approach would look like this:

from contextlib import contextmanager

def Car(name):
   print(f"__enter__ was called for {name}")
   yield name
   print(f"__exit__ was called for {name}")

with Car("toyota") as my_car:
   print(f"Driving my {my_car}")

This time we didn't have to define a whole class just to have a context manager. While this approach is much shorter to write, it also has a downside: you can't check the exception information that would be available to the __exit__() method. If you need to handle that specifically, there is no way around defining a context manager class yourself. Also note that if you need to run cleanup even on exceptions like the class-based approach does, you will need to wrap the yield statement in a try block and the cleanup code in a finally block:

def Car(name):
   print(f"__enter__ was called for {name}")
      yield name
      print(f"__exit__ was called for {name}")

So should you use context managers? It depends. In most common scenarios the benefits of with are obvious and you don't need anything else. But using it also means giving up resource cleanup tasks, meaning if you open a file or database connection using with, you don't have control over the way it is cleaned up once you are done, like a finally block would give you.

More articles

Getting help from linux man pages

Learning to navigate the linux manual pages to be productive without searching the internet

Javascript arrow functions explained

Making modern javascript code more compact

Getting more out of javascript console messages

Going beyond console.log to find more gems in the console package

Setting up an nginx load balancer

Distributing requests to multiple servers

Simplifying terminal operations with python modules

Saving time with python builtin modules

A complete guide to running windows software on linux

Making windows programs behave like they were written for linux