Python in Detail: The Context Manager Decorator

Welcome to the The Context Manager Decorator lesson!

This lesson is shown as static text below. However, it's designed to be used interactively. Click the button below to start!

In an earlier lesson, we used context managers to ensure that our files are closed, that our database transactions are committed, and that our  tags are closed with a corresponding . For that last use case, we wrote a custom context manager with the .__enter__ and .__exit__ dunder methods.
There's a surprising symmetry between context managers and generators. We'll use our Bold context manager class to see that. Here it is again.

class Bold:
  def __init__(self, file):
    self.file = file

  def __enter__(self):
    self.file.write("<b>")

  def __exit__(self, exc_type, exc_value, traceback):
    self.file.write("</b>")

Let's forget about context managers for a moment, and try to get a similar result using only the generator below. It writes "" to the output, then yields to let some other code run, then writes "".

def bold(output):
  output.write("<b>")
  yield
  output.write("</b>")

That generator is nice and short! However, the code to actually use the generator is very awkward. We recommend reading the code below carefully to understand how this generator works, despite the awkwardness. We'll improve it in a moment, but it's important to understand the details first.

Note: this code example reuses elements (variables, etc.) defined in earlier examples.

from io import StringIO

output = StringIO()
output.write("this is ")

# Call the generator to get its iterator. The generator's body hasn't started
# running yet. It's waiting for the first `next(...)` call.
iterator = bold(output)

# Iterate once. The generator writes `<b>`, then pauses at the `yield`.
next(iterator)

# Now the generator is paused and the output contains "this is <b>". We're
# "inside" the bold tag.
output.write("bold")
# Now the output contains "this is <b>bold".

# Iterate the generator again so that it writes the final `</b>`. This hits the
# end of the generator, so it raises `StopIteration`. We catch and ignore that
# exception.
try:
  next(iterator)
except StopIteration:
  pass

output.getvalue()

Result:

'this is <b>bold</b>'

There are two important things to note here. First, the bold generator function was awkward to use, but it was easy to write.
Second, we were able to rewrite a context manager as a generator with a single yield. They share the same ideas. The code before the yield is like .__enter__, and the code after the yield is like .__exit__.
It would be nice to write the simple-looking generator function, then use it as a context manager. That way, we wouldn't have to write the awkward code above.
Fortunately, that's possible. We can write a function that takes the generator function as an argument, builds a new context manager class, and returns the class. It's unusual to write a function that defines a new class, but it is possible.
Take note of the .__enter__ and .__exit__ methods here. They do the same things that we did manually in our awkward example above. In particular, the .__exit__ method has to catch and silence StopIteration in the same way.

def build_context_manager(generator):
  iterator = generator()

  class NewContextManager:
    def __enter__(self):
      next(iterator)

    def __exit__(self, exc_type, exc_value, traceback):
      try:
        next(iterator)
      except StopIteration:
        pass

  return NewContextManager

To summarize:
- build_context_manager takes a generator function.
- It gets an iterator from the generator.
- It defines a NewContextManager class.
- When we use that class as a context manager, it advances the iterator when we enter the with block, and again when we exit the with.
Now let's actually use it. In the next example, all of the awkwardness is gone. We write our context manager as a clean generator function with one yield. Then we use it in a clean with.

Note: this code example reuses elements (variables, etc.) defined in earlier examples.

from io import StringIO

def bold():
  output.write("<b>")
  yield
  output.write("</b>")

# Build the context manager class from our generator.
bold = build_context_manager(bold)
output = StringIO()
output.write("this is ")
with bold():
  output.write("bold")

output.getvalue()

Result:

'this is <b>bold</b>'

That bold = build_context_manager(bold) line sticks out. We're using a function to wrap bold. That's a job for a decorator! Let's rewrite it with @. Everything else here remains the same.

Note: this code example reuses elements (variables, etc.) defined in earlier examples.

from io import StringIO

@build_context_manager
def bold():
  output.write("<b>")
  yield
  output.write("</b>")

output = StringIO()
output.write("this is ")
with bold():
  output.write("bold")

output.getvalue()

Result:

'this is <b>bold</b>'

With build_context_manager, we can write context managers as functions, rather than writing entire classes. We wrote that decorator function from scratch, but only because it's good to see how things work. Fortunately, a more complete version of that decorator comes in the Python standard library, as contextlib.contextmanager.
Here's the same context manager defined using the @contextmanager decorator. This code is self-contained; it doesn't depend on any code from the examples above.
Here's a code problem:
The code below defines a generator function. We want to use it as a context manager. Apply the contextmanager decorator so that we can use the generator function as a context manager.
from io import StringIO
from contextlib import contextmanager
@contextmanager
def bold():
output.write("")
yield
output.write("")
output = StringIO()
output.write("this is ")
with bold():
output.write("bold")
output.getvalue()
Goal:
'this is bold'
Yours:
'this is bold'
Show Author's Answer
Let's build one more example. We want a context manager that measures the with block's runtime in milliseconds. To keep the example simple, we'll print the runtime to the console. Like the previous example, this example is self-contained; it doesn't depend on any of the code above.

from time import time
from contextlib import contextmanager

@contextmanager
def runtime():
  start = time()
  yield
  elapsed = time() - start # This gives us elapsed seconds.
  print(round(elapsed * 1000), "ms")

with runtime():
  # This loop is long enough to take a measurable amount of time.
  for _ in range(1000000):
    pass

Result:

This combines several advanced ideas: generators, iterators, context managers, and decorators. But the results are very clean. The runtime function directly says what it does: capture the current time, wait for some other code to run, then print how much time passed. When reading that function, we don't have to think about all of the details of .__enter__ and .__exit__.
There are cases where the @contextmanager decorator is a bad match. For example, remember that file objects are context managers, but they also have many other methods like .read and .write. They need to be full classes, because we need somewhere to define those methods.
Fortunately, many context managers don't need extra methods. Our bold and runtime examples above are self-contained with no methods, so @contextmanager is a good way to write them.