Context managers and the with statement

In this article we'll talk about context managers and the with keyword. What is all this for?

It's best to start with an example you already know. Say you need to open a file and write some secret — or not so secret — information into it. How would you do that, knowing how files work and about the built-in open function? Roughly like this:

PythonCopy

f = open("file.bin", "wt")
f.write("MYSECRETINFO")
f.close()

And this really will work. But now imagine you forgot to close the file and lost all the information in it. That's not very pleasant — especially if you were writing, say, your own database. Yet the point isn't even that you might forget to close the file: even if you're punctuality itself, better not to tempt fate. You wouldn't fix the wiring yourself, even knowing Ohm's law, would you? It's the same here: for working with resources it's better to reach for the language's dedicated features. That way you definitely won't shoot yourself in the foot, while achieving exactly the same result.

For working with resources it's better to reach for the language's dedicated features — that way you definitely won't shoot yourself in the foot.

What is a context manager

Let's agree on the notion of a “context manager”. A context manager is a particular structure in Python (a class or a generator) built on one main principle: when it opens and closes, code written in advance by the programmer runs.

Opening happens on entering the block via the with keyword. Closing happens when the block ends. For example: on entry a file is opened, on exit it is closed.

The simplest example of using a context manager with the with keyword is working with files:

PythonCopy

with open("file.bin", "wt") as f:  # открываем файл с помощью with
    f.write("abcdefg")

In this example we open the file via a context manager that will close it for us automatically. Why not just close the file with the .close() method? The thing is, on exiting a context manager you might need to perform dozens of operations, including calls to a third-party API or a database. Surely you don't want to write the same code every time? And even if you write a function or method that does it for you, you still have to remember to call it, otherwise it's all in vain. It's much easier to write the context manager once and call it every time we want it to do all the dirty work for us.

A class-based context manager

Context managers are very easy to write — in essence, they're just classes arranged in a special way. To write a context manager you only need to keep a few things in mind:

You need to create a class and write an __enter__ method in it. The code in this method runs on entering the context manager (when the object is created with the with keyword).

Write an __exit__ method — it runs the code placed in it on exit.

Add three extra arguments to this method besides self — exc_type, exc_val, exc_tb. We'll explain what they're for a little later.

For now, let's try writing a simple timer that measures how long our code runs:

PythonCopy

from datetime import datetime
import time  # проверять действие измерителя будем с помощью библиотеки time


# вся суть этого измерителя заключается в том, что мы считаем разницу в секундах между открытием и закрытием контекстного менеджера
class Timer:
    def __init__(self):
        pass

    def __enter__(self):  # этот метод вызывается при запуске с помощью with. Если вы хотите вернуть какой-то объект, чтобы потом работать с ним в контекстном менеджере, как в примере с файлом, то просто верните этот объект через return
        self.start = datetime.utcnow()
        return None

    def __exit__(self, exc_type, exc_val, exc_tb):  # этот метод срабатывает при выходе из контекстного менеджера
        print(f"Time passed: {(datetime.utcnow() - self.start).total_seconds()}")


with Timer():
    time.sleep(2)  # засыпаем на 2 секунды

The console should print something like the following:

TerminalCopy

Time passed: 2.00099

In addition to the comments in the code, let's say a couple of words about the arguments of the __exit__ method. Taking each of them in order:

exc_type — the type of the exception that caused the context manager to exit. If everything went fine, this argument's value will be None.

exc_val — the message in the exception. Likewise: if everything went fine, this argument will be None.

exc_tb — the traceback object from the interpreter. It's best not to touch it at all unless you're a language developer, but nevertheless it's always waiting for you here.

A generator-based context manager

Python wouldn't be Python if it didn't have a library for every last case. Back in Python 2.6 the ability to create context managers via generators appeared. Let's take a look at how this can be done.

But before that, let's clarify the difference in implementation. In essence, the only difference is that a generator-based context manager is a function (for the especially picky — a generator). Inside it, the code before yield runs the code we could have put in __enter__ if we were making the context manager as a class, and after yield we write the code that would run in __exit__. That is, before yield — everything that happens on entry, after — everything on exit. That's the whole difference; as you can see, it's small.

Now let's look at the code:

PythonCopy

from datetime import datetime
import time

from contextlib import contextmanager  # импортируем нужный нам декоратор


@contextmanager  # оборачиваем функцию в декоратор contextmanager
def timer():
    start = datetime.utcnow()
    yield  # если вам нужно что-то вернуть через контекстный менеджер, просто вставьте этот объект сюда.
    print(f"Time passed: {(datetime.utcnow() - start).total_seconds()}")


with timer():
    time.sleep(2)

The result of such a generator-based context manager is exactly the same:

TerminalCopy

Time passed: 2.001097

Summary

Let's sum up. We learned what a context manager is for and saw in practice just how convenient it is.

A context manager is invoked with the with keyword. It may or may not return an object to work with. For example, if the context manager implies working with some object, you need to add as var to the statement, where var is the variable name in this context.

We learned to build our own class-based context managers. For that we write a class with the special methods __enter__(self, ...) and __exit__(self, exc_type, exc_val, exc_tb). __enter__ is called on entering the context manager, __exit__ — on exit (don't forget to add the required arguments).

We also learned to build context managers based on generators and the contextmanager decorator from the contextlib library.

Don't be afraid to use them in your work — experiment!