Python came out with a special new keyword several years ago in Python 2.5 that is known as the with statement. This new keyword allows a developer to create context managers. But wait! What’s a context manager? They are handy constructs that allow you to set something up and tear something down automatically. For example, you might want to open a file, write a bunch of stuff to it and then close it. This is probably the classic example of a context manager. In fact, Python creates one automatically for you when you open a file using the with statement:
with open(path, 'w') as f_obj:
f_obj.write(some_data)
Back in Python 2.4, you would have to do it the old fashioned way:
f_obj = open(path, 'w') f_obj.write(some_data) f_obj.close()
The way this works under the covers is by using some of Python’s magic methods: __enter__ and __exit__. Let’s try creating your own context manager to demonstrate how this all works!
Creating a Context Manager class
Rather than rewrite Python’s open method here, you’ll create a context manager that can create a SQLite database connection and close it when it’s done. Here’s a simple example:
import sqlite3
class DataConn:
""""""
def __init__(self, db_name):
"""Constructor"""
self.db_name = db_name
def __enter__(self):
"""
Open the database connection
"""
self.conn = sqlite3.connect(self.db_name)
return self.conn
def __exit__(self, exc_type, exc_val, exc_tb):
"""
Close the connection
"""
self.conn.close()
if exc_val:
raise
if __name__ == '__main__':
db = '/home/mdriscoll/test.db'
with DataConn(db) as conn:
cursor = conn.cursor()
In the code above, you created a class that takes a path to an SQLite database file. The __enter__ method executes automatically where it creates and returns the database connection object. Now that you have that, you can create a cursor and write to the database or query it. When you exit the with statement, it causes the __exit__ method to execute and that closes the connection.
Let’s try creating a context manager using another method.
Creating a Context Manager using contextlib
Python 2.5 not only added the with statement, but it also added the contextlib module. This allows you to create a context manager using contextlib’s contextmanager function as a decorator.
Let’s try creating a context manager that opens and closes a file after all:
from contextlib import contextmanager
@contextmanager
def file_open(path):
try:
f_obj = open(path, 'w')
yield f_obj
except OSError:
print("We had an error!")
finally:
print('Closing file')
f_obj.close()
if __name__ == '__main__':
with file_open('/home/mdriscoll/test.txt') as fobj:
fobj.write('Testing context managers')
Here you import contextmanager from contextlib and decorate your file_open() function with it. This allows you to call file_open() using Python’s with statement. In your function, you open the file and then yield it out so the calling function can use it.
Once the with statement ends, control returns back to file_open() and it continues with the code following the yield statement. That causes the finally statement to execute, which closes the file. If you happen to have an OSError while working with the file, it gets caught, and the finally statement still closes the file handler.
contextlib.closing()
The contextlib module comes with some other handy utilities. The first one is the closing class which will close the thing upon the completion of the code block. The Python documentation gives an example that’s similar to the following one:
from contextlib import contextmanager
@contextmanager
def closing(db):
try:
yield db.conn()
finally:
db.close()
Basically what you’re doing is creating a closing function that’s wrapped in a contextmanager. This is the equivalent of what the closing class does. The difference is that instead of a decorator, you can use the closing class itself in your with statement.
Here’s what that will look like:
from contextlib import closing
from urllib.request import urlopen
with closing(urlopen('http://www.google.com')) as webpage:
for line in webpage:
# process the line
pass
In this example, you open a URL but wrap it with your closing class. This will cause the handle to the web page to be closed once you fall out of the with statement’s code block.
contextlib.suppress(*exceptions)
Another handy little tool is the suppress class which was added in Python 3.4. The idea behind this context manager utility is that it can suppress any number of exceptions. A common example is when you want to ignore the FileNotFoundError exception. If you were to write the following context manager, it wouldn’t work:
>>> with open('fauxfile.txt') as fobj:
for line in fobj:
print(line)
Traceback (most recent call last):
Python Shell, prompt 4, line 1
builtins.FileNotFoundError: [Errno 2] No such file or directory: 'fauxfile.txt'
This context manager doesn’t handle this exception. If you want to ignore this error, then you can do the following:
from contextlib import suppress
with suppress(FileNotFoundError):
with open('fauxfile.txt') as fobj:
for line in fobj:
print(line)
Here you import suppress and pass it the exception that you want to ignore, which in this case is the FileNotFoundError exception. If you run this code, nothing happens as the file does not exist, but an error is also not raised. It should be noted that this context manager is reentrant. This will be explained later on in this article.
contextlib.redirect_stdout / redirect_stderr
The contextlib library has a couple of neat tools for redirecting stdout and stderr that were added in Python 3.4 and 3.5 respectively. Before these tools were added, if you wanted to redirect stdout, you would do something like this:
path = '/path/to/text.txt'
with open(path, 'w') as fobj:
sys.stdout = fobj
help(sum)
With the contextlib module, you can now do the following:
from contextlib import redirect_stdout
path = '/path/to/text.txt'
with open(path, 'w') as fobj:
with redirect_stdout(fobj):
help(redirect_stdout)
In both of these examples, you are redirecting stdout to a file. When you call Python’s help(), instead of printing to stdout, it gets saved directly to the file. You could also redirect stdout to some kind of buffer or a text control type widget from a user interface toolkit like Tkinter or wxPython.
ExitStack
ExitStack is a context manager that will allow you to easily programmatically combine other context managers and cleanup functions. It sounds kind of confusing at first, so let’s take a look at an example from the Python documentation to help you understand this idea a bit better:
>>> from contextlib import ExitStack
>>> with ExitStack() as stack:
file_objects = [stack.enter_context(open(filename))
for filename in filenames]
]
This code basically creates a series of context managers inside the list comprehension. The ExitStack maintains a stack of registered callbacks that it will call in reverse order when the instance is closed, which happens when you exit the bottom of the with statement.
There are a bunch of neat examples in the Python documentation for contextlib where you can learn about topics like the following:
- Catching exceptions from __enter__ methods
- Supports a variable number of context managers
- Replacing any use of try-finally
- and much more!
You should check it out so you get a good feel for how powerful this class is.
Wrapping Up
Context managers are a lot of fun and come in handy all the time. I use them in my automated tests all the time for opening and closing dialogs, for example. Now you should be able to use some of Python’s built-in tools to create your own context managers. Be sure to take the time to read the Python documentation on contextlib as there are lots of additional information that is not covered in this chapter. Have fun and happy coding!