Write more pythonic code with context managers

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.