Python 101: Context Managers

Posted on 5 mins read

Introduction

In this post I wanted to discuss a relatively simple, but important topic: Context Managers. I want to cover the what, the why and the how.

But before we get into it... time for some self-promotion 🙊

Summary

  • Content Managers abstract away ‘clean-up’ logic.
  • Define a class with __enter__/__exit__ methods.
  • The __enter__ method is similar to a try block.
  • The __exit__ method is similar to a finally block.
  • Reduce boilerplate with @contextlib.contextmanager.

What is a Context Manager?

Officially…

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. Context managers are normally invoked using the with statement, but can also be used by directly invoking their methods. – Python Docs

In simpler terms it means: a Context Manager can ensure code that requires ‘clean-up’ logic to be executed, is done so in a more idiomatic/Pythonic way.

Why use a Context Manager?

The classic example given is when opening lots of file in Python:

files = []

for _ in range(100000):
    f = open('foo.txt', 'w')
    files.append(f)
    f.close()

Notice each file object’s .close() method is called to ensure the file descriptor is released. If we didn’t do that, then your Operating System would exhaust its allowed limit of open file descriptors.

To make this code more Pythonic and cleaner, we can utilize Context Managers.

How to use a Context Manager?

There are two ways to utilize a Context Manager…

  1. with
  2. @contextlib.contextmanager

with

We can use the with statement to define a similar block of code to our earlier ‘open multiple files’ example.

The with statement expects a ‘Context Manager’ to be provided, and there are already a few built-in Python objects designed as Context Managers; such as the open function we saw used in our above example code.

Note: another example is threading.Lock.

Here is what the code might look like when using with:

files = []

for _ in range(100000):
    with open('foo.txt', 'w') as f:
        files.append(f)

Notice how we didn’t have to explicitly call .close() on each file object generated by open. That’s because open works as a Context Manager and knows how to clean-up after itself when called via the with statement.

We’ll show you how to implement your own Context Manager in the following section: How to implement a Context Manager?.

@contextlib.contextmanager

Python provides a decorator function @contextlib.contextmanager which is actually a callable class (i.e. it defines __call__ magic method) that enables custom context managers (e.g. your own code you want to act as a context manager) to use simpler code than the traditional ‘class-based’ implementation we previously mentioned.

This means if you have custom objects that need to implement clean-up logic (similar to how open does), then you can decorate your own function so it behaves like a Context Manager, while your function itself simply uses a yield statement, like so:

from contextlib import contextmanager

files = []

@contextmanager
def open_file(path, mode): 
    file = open(path, mode)
    yield file
    file.close()

for _ in range(100000):
    with open_file('foo.txt', 'w') as f:
        files.append(f)

In the above example code we’ve effectively recreated the open Context Manager just to demonstrate the principle.

How to implement a Context Manager?

Now we’ve already seen how to implement a Context Manager using the @contextlib.contextmanager decorator (see previous sub-section), but how do we implement a class-based version of a Context Manager?

That requires us to define a class which implements __enter__ and __exit__ methods. Below is an example, again replicating the open function to keep things simple:

files = []

class Open():
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode

    def __enter__(self):
        self.open_file = open(self.filename, self.mode)
        return self.open_file

    def __exit__(self, *args):
        self.open_file.close()

for _ in range(100000):
    with Open('foo.txt', 'w') as f:
        files.append(f)

Note: for more information, see Context Manager Types.

When to use one or the other?

One thing I noticed recently was that the contextmanager variation wouldn’t execute an ’exit’ if an exception was raised during execution of the code, while the more verbose ‘class-based’ implementation would. See the following code for an example…

from contextlib import contextmanager

@contextmanager
def foo():
    print("enter!")
    yield "foobar"
    print("exit!")

try:
    with foo() as f:
        raise Exception("unexpected")
        print(f"f was: {f}")
except Exception as e:
    print(f"whoops: {e}")

The output is not what I expected:

enter!
whoops: unexpected

Notice how there is no exit! printed.

Now compare this to a ‘class-based’ example…

class Foo():
    def __enter__(self):
        print("enter!")

    def __exit__(self, *args):
        print("exit!", args)

try:
    with Foo() as f:
        raise Exception("unexpected")
        print(f"f was: {f}")
except Exception as e:
    print(f"whoops: {e}")

The output is as expected:

enter!
exit! (<class 'Exception'>, Exception('unexpected'), <traceback object at 0x108882d00>)
whoops: unexpected

i.e. we see both an enter and exit message.

We can get the contextmanager to behave as we might have expected it to (e.g. the same as the ‘class-based’ implementation) by ensuring the function that calls yield is wrapped in a try/finally block, like so:

from contextlib import contextmanager

@contextmanager
def foo():
    print("enter!")
    try:
        yield "foobar"
    finally:
        print("exit!")

try:
    with foo() as f:
        raise Exception("unexpected")
        print(f"f was: {f}")
except Exception as e:
    print(f"whoops: {e}")

The output of this is now what we might expect…

enter!
exit!
whoops: unexpected

When choosing between the two options contextmanager and ‘class-based’ implementation, it might be worth keeping this caveat in mind.

Multiple Context Managers in a single With statement

One interesting aspect of the with statement is that you can execute multiple context managers as part of its block control. Meaning when the with block completes, then all context managers will be cleaned up.

from contextlib import contextmanager

@contextmanager
def foo():
    print("enter!")
    try:
        yield "foobar"
    finally:
        print("exit!")


with foo() as f1, foo() as f2, foo() as f3:
    print(f"f1 was: {f1}")
    print(f"f2 was: {f2}")
    print(f"f3 was: {f3}")

Alternatively you can utilize contextlib.ExitStack:

from contextlib import contextmanager, ExitStack

@contextmanager
def foo():
    print("enter!")
    try:
        yield "foobar"
    finally:
        print("exit!")

with ExitStack() as stack:
    managers = [stack.enter_context(foo()) for cm in range(3)]
    print(managers)  # ['foobar', 'foobar', 'foobar']

But before we wrap up... time (once again) for some self-promotion 🙊