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
try:
f = open("test.txt")
print(f.read())
finally:
if f is not None:
f.close()
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:
print(f.read())
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:
print(f1.read())
print(f2.read())
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
@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:
@contextmanager
def Car(name):
print(f"__enter__ was called for {name}")
try:
yield name
finally:
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.