Using context managers in Python

Using context managers in Python

If you ever used Python, most likely you used context managers without even knowing what they are and what they do or how they work - you just used them as per the documentation of some 3rd party package. In this article, I will try to explain what exactly are context managers and what are their benefits.

1. Context managers (by definition)

I hate to be a bookworm (I really do, trust me), but sometimes good old definitions come in handy. This is the case with context managers. So here it goes:

Context managers are responsible for allocating and releasing resources when needed.

Is a question mark popping up in your head? Expected. I will walk you with a practical example :)

1.1 Practical example - no context manager

The idea is to open up a file with Python and read the file contents of that specific file. Let's take a look at the code below:

file = open("textfile.txt", "r")
file.read()
file.close()

Simple eh? We open up a file called textfile.txt in read mode (r) and we read all of the contents of that file by making use of the read() method. Once we are done with injecting the content inside of the memory (or a variable if you want), we close the file.

Now - imagine that you forgot to close the file. Now imagine that you work with tons of files and somehow forget to close all of them. This can lead to substantial resource leaks and you want to avoid that.

Enter context managers :)

1.2 Practical example - context manager

The idea is once again to open up a file and read its contents. This time we will use context manager. The whole idea of a context manager is to do the following:

  • Open up the resource (allocate it)
  • Do something with the resource
  • Close the resource once it is not needed anymore (release it)

In Python, we use the context managers by making use of the with keyword. Let's see what the code looks like now:

with open("textfile.txt", "r") as file:
    file.read()
    # some other code while the resource (file) is allocated

Similar, yet different. Where did the file.close() go? Funny that you ask - we don't need it anymore.

But... but... you said that we need to close the file to avoid resource leaks bla bla...

Yes, I did, and it is true. But what if I told you that context managers do that for you. In the above code, the resource (file) is opened as a variable called file in a read mode and is made available for usage (additional related code) inside of a block (indentation). Inside of that block, we read the file content. We can also do additional stuff with the file INSIDE of that block, BUT once we leave that block, the context manager will release the resource (in this case close the file) and we can no longer access it.

Handy eh? Let's see how it works under the hood.

2. Context manager under the hood

Each object (in Python everything is an object) that wants to use context manager, by using the with keyword, has to implement two built-in methods - __enter__ and __exit__.

For the above use case, if we were to create a custom context manager class it would look something like this:

from types import TracebackType
from typing import IO, Optional, Type


class CustomClassFile:
    def __init__(self, file_path: str, file_mode: str) -> None:
        self.file = open(file_path, file_mode)

    def __enter__(self) -> IO:
        return self.file

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_value: Optional[BaseException],
        traceback: Optional[TracebackType],
    ) -> None
        self.file.close()


with CustomClassFile("textfile.txt", "r") as file:
    file.read()

So, once we create a new instance of CustomClassFile, we pass the two arguments - file_path and file_mode. That is fairly simple.

Inside of the __enter__ method is where we need to define the resource that we want to allocate once the context manager kicks in. In this case, it will be the self.file object.

The __exit__ method signature is a bit weird, but I won't bother you with the details of it. All of those arguments are required by default - so whenever you will create the __exit__ method, you would use this exact same signature. This is due to the ability to create your own exception handlers (if they occur) upon releasing the resource. What is important is that we call the self.file.close() method inside of the __exit__ method. This means that once we get past the with block, the file will be closed automatically by the context manager :)

Let's see some more examples.

2.1 Context manager as a class by example

The idea here would be to create the custom class that will allocate not the file but rather the contents of a JSON file and releases the file afterward.

import json
from types import TracebackType
from typing import Any, IO, Optional, Type


class JSONFileLoader:
    def __init__(self, file_path: str) -> None:
        self.file: IO = open(file_path, "r")
        self.data: Any = json.load(self.file)

    def __enter__(self) -> Any:
        return self.data

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_value: Optional[BaseException],
        traceback: Optional[TracebackType],
    ) -> None:
        if self.file:
            self.file.close()


with JSONFileLoader("file.json") as json_data:
    print(json_data)

This time when we create a new instance of the JSONFileLoader we only pass one argument - file_path. The __init__ method then also loads the json contents of the file inside of the self.data property.

The __enter__ method will this time allocate the data (as per the idea above) instead of the file itself. That means that once we enter the context manager instead of the whole file like before we will only have the file data available.

The __exit__ method remains the same as before - we just close the file.

with JSONFileLoader("file.json") as json_data:
    print(json_data)

See that now instead of the file, our __enter__ method actually is returning the self.data property, which means that json_data is actually whatever is set inside of the self.data property. This makes it simple to just - for this example - print out the data. Once we leave the with block, __exit__ will kick in and close the file for us.

2.2 Context manager as a generator by example

This is an alternative approach to creating context managers in Python. If you don't know what generators are in Python, please take a moment to read about them.

Let's see how to make use of Python's contextlib module to solve the same idea as the one above.

import json
from contextlib import contextmanager
from typing import Any, Generator, IO


@contextmanager
def json_file_loader(file_path: str) -> Generator[Any, None, None]:
    file: IO = open(file_path, "r")
    data: Any = json.load(file)
    yield data
    file.close()


with json_file_loader("file.json") as json_data:
    print(json_data)

So the idea is that you will create a function and decorate it by using @contextmanager to mark it as a context manager for Python. We create a function that takes the file_path as an argument and it opens the file and loads it by using the json.load() method. Once that is done, we yield the data. That means we will pause the execution of the code below from the function json_file_loader until we exit the context manager. This means that file.close() will execute ONLY after we exit the context manager as described below.

with json_file_loader("file.json") as json_data:
    # inside of the context manager (yield kicked in)
    print(json_data)
# context manager exited which means that file.close() will execute

3. Benefits of using context managers

So what would be the perks of using context managers?

  • Reduce resource leaks
  • Reduce the number of lines of code
  • More power over what to expose inside of the with block
  • Multiple ways to implement
  • No need to worry about releasing the resources

I hope you find those benefits handy and start to make use of context managers.

Like always, thanks for reading :)