One of the first things you’ve learnt when handling files in Python was definitely with statements. Everyone taught you to use it, so you don’t have to remember about allocating resources and closing your own files. And since then you probably always used this magical:

with open("file.txt", "r") as f:
    f.read()

But what if I told you that you can actually implement your own contexts to use alongside those with statements? Wouldn’t it make your code shorter, prettier and more readable in many cases?

In this article I will:

• explain to you what are context managers,
• show you step by step what happens under the hood,
• show you how to implement your own context managers.

Let’s go!

What is a context manager?

If we look up a definition of context manager in Python docs, we’ll find that:

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.

Uhm, right… Aren’t the overly complex definitions one of the reasons of bad communication in Data Science? Let’s decompose this definition to understand precisely what we’re dealing with.

A context manager is an object…

Ok so this part actually tells us that we’re dealing with an instance of some Python class. Now what this class will represent?

…that defines the runtime context…

This means our class will represent some type of environment (context) where things will be managed for us automatically (manager). 

…to be established when executing a with statement.

So this will start being managed once we use our context manager alongside the with statement. Now how are things going to be managed?

The context manager handles the entry into, and the exit from, the desired runtime context…

Meaning – something will happen when we enter the context and then again when we are going to leave it, so it needs to have enter and exit methods.

That tells us a lot! Let’s see our file opening example again to see if our understanding makes sense.

with open("file.txt", "r") as f:
    f.read()

What roughly happens here under the hood is:

• we enter the context by opening a file and storing it in variable f,
• we try executing whatever is in our body – in this case f.read(),
• we exit the context by closing a file.

It appears then that such neat mechanism allows us to:

• automate operations that we have to do every time before and after some chunk of code,
• and therefore reduce the unnecessary code in our script.

Let’s see a couple of machine learning contexts’ examples and how to implement them from scratch.

Implementation of custom context manager as a class

In the previous section, we realised that context manager is actually a class. Let’s see an example of such an implementation for switching between training and eval mode in PyTorch.

class TrainingSession:

    def __init__(self, model: torch.nn.Module):
        self.model = model

    def __enter__(self):
        self.model.train()
        return None

    def __exit__(self, exc_type, exc_value, exc_traceback):
        self.model.eval()

Similarly to our file example, here we also have enter and exit methods. Upon entering our TrainingSession, we will switch the model into training mode, and then switch to evaluation after exiting it. See for yourself.

with TrainingSession(model=model):
    print(model.training)            # -----> this results in True
print(model.training)                # -----> this results in False

There are a couple of additional things about context managers we can talk about in this particular example.

Operator: as

You definitely remember from file handling example that we were able to save result in a variable

with ... as something:
    ...

In our example if we did:

with TrainingSession(model=model) as session:
    ...

a session will be equal to whatever return in enter gives (in this case None). We can change it and add some more parameters that will be accessible for us in context body:

class TrainingSession:

    def __init__(self, model: torch.nn.Module, epochs: int):
        self.model = model
        self.epochs = epochs

    def __enter__(self):
        self.model.train()
        return {"epochs": self.epochs}

    def __exit__(self, exc_type, exc_value, exc_traceback):
        self.model.eval()

We added epochs argument to init and changed enter to return dict. Now we will be able to access some of the Training Session parameters we defined in our body:

with TrainingSession(model=model, epochs=2) as session:
    for i in range(session["epochs"]):
        print(i)

This returns

0
1

exactly as expected. Nice!

Input to exit

As you can see in the TrainingSession example, exit method takes as input some mysteriously looking arguments like exc_type, exc_value and exc_traceback. Let’s analyse it.

When we defined step by step what happens inside the context, we said that we try executing the code in the body and not that we execute it. Why? Because code in the body of the context can actually fail! But no matter if it fails or not, we would still like to execute our exit function properly.

If everything happens as expected, all those 3 arguments are set to None. Now let’s see what happens when the code fails. Firstly, I add to exit to print out our arguments and to always return True:

class TrainingSession:

    ...

    def __exit__(self, exc_type, exc_value, exc_traceback):
        print(exc_type, exc_value, exc_traceback)
        self.model.eval()
        return True

and then something that should fail (we will try addressing non existent unknown_arg in our session):

with TrainingSession(model=model, epochs=2) as session:
    print(session["unknown_arg"])
print(model.training)

The output is:

<class 'KeyError'> 'unknown_arg' <traceback object at 0x7f4fbf067aa0>
False

First of all, because we set return True in exit, our code did not raise error and fail – it now relies on us to handle the problem ourselves. And in order to handle the problem, we need to know what type of problem happened – this is why we get those 3 arguments and we can see:

• exc_type – what type of error we got (KeyError),
• exc_value – what was the error message ('unknown_arg'),
• exc_traceback – traceback object (<traceback object at 0x7f4fbf067aa0>).

And one important thing is that we still fully executed exit – meaning our model returned to eval mode after exiting TrainingSession, even if the body failed.

Summary

So what should you remember from this article? Let’s see:

• context managers allows us to create “spaces” for code execution and manage what happens once we enter and exit this space,
• they are executed using with statement,
• in enter, we can represent context manager as another object that will be available in our space (see how we started treating session as dict of parameters),
• in exit, we can handle ourselves potential failures of our code and still execute important parts (like switching model mode).

Thanks for reading! I’m proud of you for reaching here!