Write more pythonic code with context managers
![](/assets/img/posts//write-more-pythonic-code-with-context-managers/write-more-pythonic-code-with-context-managers-cover.png)
Batteries included is a blog series about the Python Standard Library. Each day, I share insights, ideas and examples for different parts of the library. Blaugust is an annual blogging festival in August where the goal is to write a blog post every day of the month.
In Python community, we often talk about writing pythonic code. What that exactly means is up for interpretation but in essence, it’s about writing Python code that takes advantage of Python’s features and language constructs in the best way. Sometimes, it’s used in the context of discussing PEP 8 which is the Python style guide but that’s a rather surface level description.
Long-time Python core developer Raymond Hettinger has a wonderful talk about pythonic code from PyCon US 2015 called Beyond PEP 8 which is highly recommended watching for everyone who writes Python.
Today’s post is about Context Managers, a construct in Python language that makes the code easier to read and understand, making it more pythonic. To quote the documentation:
A context manager is an object that defines
the runtime context to be established when executing a
with
statement. The context manager handles the entry into, and the exit from, the
desired runtime context for the execution of the block of code.
If you have been writing Python somewhat recently, the most likely case where you have encountered them is in the context of opening files:
with open('filename.txt', 'r') as datafile:
for line in datafile:
print(line)
There’s nothing preventing you from achieving the same without context managers:
datafile = open('filename.txt', 'r')
for line in datafile:
print(line)
datafile.close()
but I would argue that the context manager makes it clearer to read and understand. And it’s less likely you forget to close the reader after you’re done with it as this context manager closes the file at the end of the block.
Any class can be made into a context manager by defining two methods:
__enter__
and
__exit__
. When the context manager is
entered in the with
block, the
__enter__
method is run and when the
block is finished, the __exit__
method
is run.
In the aforementioned Raymond Hettinger talk, he has a great example of making code more pythonic by wrapping networking code initializing and finalizing into a context manager (note that this is Python 2 code):
from nettools import NetworkElement
with NetworkElement('171.0.2.45') as ne:
for route in ne.routing_table:
print '%15s -> %s' % (route.name, route.ipaddr)
with the NetworkElement
defined as (only
showing the __enter__
and
__exit__
)
class NetworkElement(object)
def __enter__(self):
return self
def __exit__(self, exctype, excinst, exctb):
if exctype == NetworkElementError:
logging.exception('No routing table found')
self.oldne.cleanup('rollback')
else:
self.oldne.cleanup('commit')
self.oldne.disconnect()
You can do a lot of validation, cleaning up, disconnecting and error handling inside the context manager, leaving the user of the context manager to focus on the main logic they want to implement.
Another handy use case for context managers is with testing exceptions with pytest:
import pytest
def test_zero_division():
with pytest.raises(ZeroDivisionError):
1 / 0
You can also use context managers when mocking tests with patch:
from unittest.mock import patch
with patch('Class.method', mocked_method):
# Run your tests
# once the block is over, the original method is put back
Trey Hunter has a great series of screencasts for context managers if you want to learn more.