Chapter 3 - Context Managers
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:
1 with open(path, 'w') as f_obj:
2 f_obj.write(some_data)
Back in Python 2.4, you would have to do it the old fashioned way:
1 f_obj = open(path, 'w')
2 f_obj.write(some_data)
3 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 our own context manager to demonstrate how this all works!
Creating a Context Manager class
Rather than rewrite Python’s open method here, we’ll create a context manager that can create a SQLite database connection and close it when it’s done. Here’s a simple example:
1 import sqlite3
2
3
4 class DataConn:
5 """"""
6
7 def __init__(self, db_name):
8 """Constructor"""
9 self.db_name = db_name
10
11 def __enter__(self):
12 """
13 Open the database connection
14 """
15 self.conn = sqlite3.connect(self.db_name)
16 return self.conn
17
18 def __exit__(self, exc_type, exc_val, exc_tb):
19 """
20 Close the connection
21 """
22 self.conn.close()
23 if exc_val:
24 raise
25
26 if __name__ == '__main__':
27 db = '/home/mdriscoll/test.db'
28 with DataConn(db) as conn:
29 cursor = conn.cursor()
In the code above, we created a class that takes a path to a SQLite database file. The __enter__ method executes automatically where it creates and returns the database connection object. Now that we have that, we can create a cursor and write to the database or query it. When we 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 us 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:
1 from contextlib import contextmanager
2
3 @contextmanager
4 def file_open(path):
5 try:
6 f_obj = open(path, 'w')
7 yield f_obj
8 except OSError:
9 print("We had an error!")
10 finally:
11 print('Closing file')
12 f_obj.close()
13
14 if __name__ == '__main__':
15 with file_open('/home/mdriscoll/test.txt') as fobj:
16 fobj.write('Testing context managers')
Here we just import contextmanager from contextlib and decorate our file_open function with it. This allows us to call file_open using Python’s with statement. In our function, we open the file and then yield it out so the calling function can use it.
Once the with statement ends, control returns back to the file_open function and it continues with the code following the yield statement. That causes the finally statement to execute, which closes the file. If we happen to have an OSError while working with the file, it gets caught and finally statement still closes the file handler.
contextlib.closing(thing)
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 code block. The Python documentation gives an example that’s similar to the following one:
1 from contextlib import contextmanager
2
3 @contextmanager
4 def closing(db):
5 try:
6 yield db.conn()
7 finally:
8 db.close()
Basically what we’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, we can use the closing class itself in our with statement. Let’s take a look:
1 from contextlib import closing
2 from urllib.request import urlopen
3
4 with closing(urlopen('http://www.google.com')) as webpage:
5 for line in webpage:
6 # process the line
7 pass
In this example, we open a url page but wrap it with our closing class. This will cause the handle to the web page to be closed once we 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. Let’s say we want to ignore the FileNotFoundError exception. If you were to write the following context manager, it wouldn’t work:
As you can see, this context manager doesn’t handle this exception. If you want to ignore this error, then you can do the following:
1 from contextlib import suppress
2
3 with suppress(FileNotFoundError):
4 with open('fauxfile.txt') as fobj:
5 for line in fobj:
6 print(line)
Here we import suppress and pass it the exception that we want to ignore, which in this case is the FileNotFoundError exception. If you run this code, you will note that 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 section.
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:
1 path = '/path/to/text.txt'
2
3 with open(path, 'w') as fobj:
4 sys.stdout = fobj
5 help(sum)
With the contextlib module, you can now do the following:
1 from contextlib import redirect_stdout
2
3 path = '/path/to/text.txt'
4 with open(path, 'w') as fobj:
5 with redirect_stdout(fobj):
6 help(redirect_stdout)
In both of these examples, we are redirecting stdout to a file. When we 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 us 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 it closed, which happens when we exit the 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!
I highly recommend checking it out so you get a good feel for how powerful this class is.
Reentrant Context Managers
Most context managers that you create will be written such that they can only be used once using a with statement. Here’s a simple example:
>>> from contextlib import contextmanager
>>> @contextmanager
... def single():
... print('Yielding')
... yield
... print('Exiting context manager')
>>> context = single()
>>> with context:
... pass
...
Yielding
Exiting context manager
>>> with context:
... pass
...
Traceback (most recent call last):
Python Shell, prompt 9, line 1
File "/usr/local/lib/python3.5/contextlib.py", line 61, in __enter__
raise RuntimeError("generator didn't yield") from None
builtins.RuntimeError: generator didn't yield
Here we create an instance of our context manager and try running it twice with Python’s with statement. The second time it runs, it raises a RuntimeError.
But what if we wanted to be able to run the context manager twice? Well we’d need to use one that is “reentrant”. Let’s use the redirect_stdout context manager that we used before!
>>> from contextlib import redirect_stdout
>>> from io import StringIO
>>> stream = StringIO()
>>> write_to_stream = redirect_stdout(stream)
>>> with write_to_stream:
... print('Write something to the stream')
... with write_to_stream:
... print('Write something else to stream')
...
>>> print(stream.getvalue())
Write something to the stream
Write something else to stream
Here we create a nested context manager where they both write to StringIO, which is an in-memory text stream. The reason this works instead of raising a RuntimeError like before is that redirect_stdout is reentrant and allows us to call it twice. Of course, a real world example would be much more complex with more functions calling each other. Please also note that reentrant context managers are not necessarily thread-safe. Read the documentation before trying to use them in a thread.
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!