Context manager in Python

Today we're gonna talk about the "with" statement and the context managers in Python.

The "with" statement relates to something that is called context managers. And it is sometimes regarded as a bit of an obscure feature by some people. But actually when you peek behind the scenes there is relatively little magic involved. And it's actually a highly useful feature that can help you write cleaner and easier to understand, easier to maintain code.

So the first question that I wanna talk about is what is the "with" statement actually good for?

The "with" statement is great for simplifying some of the common resource management patterns that you'll see in your code. By resources in this case I mean system resources, like files, locks, or some kind of network connections, stuff like that. And so with the "with" statement you can abstract the way some of the functionality there that is commonly required, like dealing with the acquisition and releasing these resources again. A good example here would be dealing with files.

An open file is a system resource. And you need to acquire it in some way and also release it to give that resource back to the system, so it doesn't have to keep track of the file you're working with.

Therefore I wanted to look at this simple example here, where we're using the "with" statement to open a file and then write to it, and to have the file automatically closed.

with open('file.txt', 'w') as f:
    f.write('CheckiO is the best coding game')

The big advantage here is that if we use the "with" statement to open the file, Python is gonna make sure that this file is gonna be closed once we're done with it. Basically, when you do this - with open('file.txt', 'w') as f:, you open or you create this context where you're working with the file object, we're giving it a name "f". And then in the block that follows, in this indented block here, we can work with f, and as soon as the execution of your program leaves the context, Python will automatically close the file. Ergo it's gonna call f.close and that will return the resource, the file descriptor, back to the operating system and it doesn't have to keep track of that resource anymore.

To demystify how the "with" statement works let's take a look at what actually happens behind the scenes. In the event you were to write this manually, you'd probably do something like that. You would open the file, just assigned it to f here, on this line, and then I'm using a try..finally statement.

f = open('file.txt', 'w')
try:
    f.write('CheckiO is the best coding game')
finally:
    f.close()

Trying here to write to the file, and then to make sure I actually close the file. So, even if something goes wrong here, we have some kind of exception, some error and we can't write to the file, try..finally would make sure that we're actually calling close on the file, which in a longer running program that's super important because you always wanna give up these resources. Consequently the try..finally here is really significant. It wouldn't be enough to just say: "OK, I'm gonna do an open, and then write, and then close", because you could potentially leak that resource by not releasing it again if some kind of exception would happen.

Using the "with" statement really simplifies the code we have to write here. Because what you see in try-finally code example, that is super common. You are gonna use a pattern like this very very frequently if you're dealing with any kind of resources like open files, open network connections, stuff like that. You always wanna make sure that you're using this try..finally pattern. Thus the "with" statement is a way to abstract that away, to factor out that functionality, so that you don't have to write that every single time you open a file, you create a file.

Now, how does this actually work behind the scenes? Because this seems like a really nice feature. In my opinion, it's a highly useful feature in Python. How can you support the "with" statement in your own objects? Because there's really nothing magical or special about the way open works. There's not some magic sauce that's only available to Python builtin objects. You can support the same functionality in your own programs or in your own objects by implementing so-called context managers.

What's a context manager? Really, what it boils down to is that it's a simple protocol or some interfacer contract that your objects follow, so that they can be used with the "with" statement. And I'm gonna show you a simple example here.

class ContextFile:
    def __init__(self, file_name):
        self.file_name = file_name

    def __enter__(self):
        self.file = open(self.file_name, 'w')
        return self.file

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.file:
            self.file.close()

As you can see this ContextFile class basically emulates what open did there and how we were able to use it with the "with" statement. There are two methods that an object needs to support in order to be used with the "with" statement, and that is the __enter__ and the __exit__ method. Because it's a very very simple contract and a very very simple interface. And with this ContextFile class here we can go ahead and say:

with ContextFile('hello.txt') as f:
    f.write('hello, world!')
    f.write('bye now')

that looks exactly the same as the open function called, looked earlier, and now we can actually go ahead and write something to that file, and we gonna get the same result, the file is automatically gonna be closed.

Now I wanna talk a little bit more about the steps that Python takes behind the scenes for this example to actually work. You can see here we defined this ContextFile class and we've got this constructor __init__ that remembers the name for the file it wanna create, and it doesn't actually open the file until the __enter__ method is called. And that's kind of how you wanna structure your context managers, so that really the resource gets acquired when the __enter__ method is called, __enter__ returns the resource as a result and it gets released when the __exit__ method is called. And then the __exit__ method also takes additional parameters that will tell you about some exception that might have happened in case you wanna inspect that log or do something with that.

Obviously this is sort of a useless reper around the open function here because the open function already pretty much does that when it functions as a context manager. This is merely an example to illustrate that. But I hope that gives you a better idea of how these context managers actually work behind the scenes.

Now, I've just explained to you how class based context managers work. But this isn't the only way to support the "with" statement in Python and this is not the only way to implement a context manager. There is the contextlib module in a standard library and it provides a couple of abstractions on top of the basic context manager protocol.

Here's a quick example of what you can do with this. I'm gonna reimplement the same ContextFile functionality using the contextlib library.

from contextlib import contextmanager

@contextmanager
def context_file(name):
    try:
        f = open(name, 'w')
        yield f
    finally:
        f.close()

There's a decorator in there called contextmanager and this thing is highly useful. What that allows you to do is you can define objects that follow the context manager protocol that you can use with the "with" statement simply by writing a generator. Looking at this code again you can see our familiar try..finally pattern here where I'm acquiring the resource and then I'm yielding it, and then later I'm closing the resource. And what happens is that this @contextmanager decorator will turn this generator function or this generator that I just defined here, it'll turn that into a full blown context manager that I can use with the "with" statement.

With this technique using contextlib and using the contextmanager decorator you can write some of this context managers a lot quicker. I guess, the downside is that for someone to understand this piece of code they would have to have some basic knowledge about decorators and they would have to have some basic knowledge about how generators work in Python.

I hope that I gave you a good overview of how these context managers work, how the "with" statement works behind the scenes and what you can do with them.