Tidy Up After Yourself: Python Context Managers

Resource leaks are a problem that will sneak up on you during the execution of your program. Like their more famous relative, memory leaks, they're caused when you allocate something and then don't release it when you've finished with it.

Sometimes their effect is obvious. Your system might start and then crash immediately. That's good, because you're obviously calling the guilty code a lot, and that means you have a better idea where to look for the problem. Things get more complicated when the resource is only used occasionally.

Why Am I Getting Resource Leaks?

Whenever you create an object in your code, it needs to be deleted at some stage. Some languages need you to do this explicitly, e.g. the delete statement in C++. Others handle it automatically with garbage collection. This is great when you're using objects that CAN be deleted automatically – things like data structures etc. Some data types can't be deleted quite so simply though. Things that use external resources like file handles, database connections and communications links can't be trusted to the language run-time. You have to handle the tidying-up yourself.

The most basic way to do this is some variation on the following piece of C# code:

MyResourceHandlingClass myItem = null;

try
{
    MyResourceHandlingClass myItem = new MyResourceHandlingClass(a, b, c));

    // Do things with myItem

}
catch ...
{
    // Handle exceptions
}
finally
{
    // Tidy up any member objects
    if (myItem != null)
    {
        myItem.ReleaseResources();  // Custom tidy-up method
        myItem = null; // Optional tidying up
    }
}

This is a lot of work to ensure just one object can't leak resources. Imagine how much worse it gets when you have two or three objects in use at the same time.

Thankfully, plenty of languages let you encapsulate the creation and destruction of objects in some form of context. It's a handy piece of syntactic sugar to cover up the cleanup of an object regardless of whether any exceptions were thrown. When your execution is inside this context, you don't need to write the extra clean-up code; it's all done for you when you leave the context scope.

Doing this in C# is simple - you just implement your class's cleanup code using the IDisposable interface and pattern. Then you use a using block. For example:

using (MyDisposableClass myItem = new MyDisposableClass(a, b, c))
{
    // Do things with myItem
}
// At this point, myItem.Dispose() has been called

Python, as always, is a little different. You've almost certainly seen code like this in Python:

with open("myfile.txt", mode="rw+") as my_file:
    # Do things with my_file

# my_file is closed by the time you get here

This construct ensures that myfile.txt is flushed and closed, assuming it was created properly. If an exception occurs then, when the scope is closed, things are still cleaned up properly. This means that if the file was not created properly, it does not exist. If it was created properly, then the contents depend on how much data was written before the exception, but file is still closed though.

This All Looks Rather Good! How Can I Do This?

The Python context manager can be written as a class with three methods, like this:

class MyContextManager:
    def __init__(self, param1, param2):
        self._first = param1
        self._second = param2
      
    def __enter__(self):
        # Create managed object(s) here, and store and/or return them
        return <an object>

    def __exit__(self, exc_type, exc_value, traceback):
        # Do tidying up here, and handle any exception(s) that arrive

It's important to remember that the context manager is an object that can contain data members. Normally you would return the self object from __enter__, and then call functions on that. You can also return a member of the context object. This is how you use the file object in the following example, where __enter__ is returning a file object:

with os.open("myfile.txt", "r", encoding="utf-8") as my_file:
    contents = my_file.readlines()

The definition of the os.open context manager looks something like:

class open:
    def __init__(self, file_path, mode, encoding):
        self._file_path = file_path
        self._mode = mode
        self._encoding = encoding
      
    def __enter__(self):
        self._file = open(self._file_path, self._mode, self._encoding)
        return self._file

    def __exit__(self, exc_type, exc_value, traceback):
    
        # Do something with the exception

        self._file.close()
        self._file = None

But you don't HAVE to do this. The context manager is just a regular object, so it can contain data members. That means you can refer to them in the __enter__ and __exit__ methods, and return something else from __enter__ if it suits you. For example, the tempfile.TemporaryDirectory class returns the path to the temporary directory it creates from the __enter__ method, and looks after the actual directory management within the context manager class itself. This means that you can create files in the temporary directory and know that when you code leaves the context manager, all of the files and the temporary directory will be deleted without you having to worry about them. It saves you a LOT of work keeping track of files, checking whether they exist, deleting them and removing the directory.

Most of the examples so far have featured the file system, but you can also use context managers to look after network connections, database connections or anything else you want precise control over.

Naturally, context managers can be used in asynchronous code, just use the following pattern if you want to create your own class:

class MyAsyncContextManager:
    def __init__(self, param1, param2):
        self._first = param1
        self._second = param2
      
    async def __aenter__(self):
        # Create managed object(s) here, and store and/or return them
        return await <async created object>

    async def __aexit__(self, exc_type, exc_value, traceback):
        # Do tidying up here, and handle any exception(s) that arrive

Then use this context manager like:

async with MyAsyncContextManager() as manager:
    await manager.function()

Writing My Own Context Managers Seems Like Too Much Work! Are There Any Short Cuts?

Of course there are - the contextlib module contains lots of short cuts! Some of the most useful:

If you have a function that returns an object you want managed by a context manager, you can use the @contextmanager decorator to turn a generator-iterator function that returns a managed object with a yield statement, and then deletes it afterwards. There is, of course, an asynchronous version @asynccontextmanager in case you want that.

If you have a class that has a synchronous close() method (or asynchronous aclose() method) to tidy up after itself, use the closing() or asynchronous aclosing()) context manager to automatically call the appropriate method.

The official Python documentation even has examples and recipes for you to take inspiration from.

Once you're familiar with them, you'll be using them in places you wouldn't have imagined. Just remember, the result of the __enter__() or __aenter__() method is usually the context manager itself, but it doesn't have to be. Take some inspiration from the tempfile.TemporaryDirectory class and go wild with it!

Summary

Context managers are an excellent way to simplify your code AND to reduce the risk of resource leaks. These will make your code easier to read, and will improve the reliability of your systems. Both will be appreciated by your colleagues and users. Try them, and you'll be impressed.