Context Manager and Contextlib in Python

Let's dig into the context manager and contextlib in Python and see how it can be created.

We'll cover the following

About 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 us to set something up and tear something down automatically. For example, we 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 us when we open a file using the with statement:

Press + to interact
with open(path, 'w') as f_obj:
    f_obj.write(some_data)

Back in Python 2.4, we would have to do it the old fashioned way:

Press + to interact
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:

Press + to interact
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 = "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 (Line #11). 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 (Line #18).

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

Creating a context manager using contextlib

Let’s try creating a context manager that opens and closes a file after all:

Press + to interact
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('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.