Python 201 – An Intro to 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:

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 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:

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 __name__ == '__main__':
    db = '/home/mdriscoll/test.db'
    with DataConn(db) as conn:
        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:

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 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.

---

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. Have fun and happy coding!

---

Related Reading

• contextlib documentation
• Understanding Python’s “with” statement