Python in Detail: Custom Context Managers

Welcome to the Custom Context Managers lesson!

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

We've seen that context managers automatically perform cleanup operations, like closing files when we're done with them.

with open("data.txt") as file:
  data = file.read()

Calling open("data.txt") gives us a file object. It supports all of the usual file operations, but it's also a context manager. The as file puts the file in a variable.
The with statement doesn't know anything about files (or database transactions, or any other kind of resource management problem). Instead, the object that we're managing needs to define two dunder methods: .__enter__ and .__exit__. That object is the x in with x. For example, if we're managing a file via with open("data.txt"), the file returned by open must define the methods.

from io import StringIO

# We use a StringIO instead of a real file. They're identical for the purposes
# of this example.
file = StringIO()

(hasattr(file, "__enter__"), hasattr(file, "__exit__"))

Result:

Using an object in a with statement does two special things. When execution enters the with, Python calls .__enter__ on the object. That gives the context manager (like the file) a chance to do any setup that it needs. Then all of the code inside of the with executes as usual.
After the last statement in the with, Python calls .__exit__ on the object. That gives the context manager a chance to do any cleanup that it needs. In the case of files, this is when the file closes itself.
Many context manager use cases involve external resources, like files or network connections or databases. In the next few examples, we'll build a context manager that doesn't need any external resources.
Imagine that we're writing some HTML data to a StringIO, and we want to make sure that HTML tags are always closed. For example, bold text should always be wrapped in a  at one end and a  at the other end. We'll use a context manager to make sure that we never forget to close the tag.

from io import StringIO

class Bold:
  # Our examples use `StringIO`, but this class will work with any file-like
  # object. We name the argument "file" to make that clear.
  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>")

When the with statement starts, Python calls .__enter__, which outputs "". When the with ends, Python calls .__exit__, which outputs "".

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

output = StringIO()

output.write("an ")
with Bold(output):
  output.write("important")
output.write(" thing")
output.getvalue()

Result:

Here's a code problem:
Write a Paragraph context manager. It should wrap any text in an HTML paragraph ....
from io import StringIO

file = StringIO()
class Paragraph:
def __init__(self, file):
self.file = file

def __enter__(self):
self.file.write("")

def __exit__(self, exc_type, exc_value, traceback):
self.file.write("")
with Paragraph(file):
file.write("hello world")

file.getvalue()
Goal:
'hello world'
Yours:
'hello world'
Show Author's Answer
What if the code in the with block raises an exception? Python still calls the context manager's .__exit__ method. This is similar to finally: clauses when catching exceptions: they always run, whether an exception is raised or not.

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

output = StringIO()

output.write("an ")

try:
  with Bold(output):
    output.write("exception")
    raise ValueError()
except ValueError:
  pass

output.write(" caught")
output.getvalue()

Result:

'an <b>exception</b> caught'

You may notice that our .__exit__ method above took three arguments, but we didn't use them. When no exception happens, those arguments are all None. But when the with block raises an exception, Python passes values to all three arguments:
- exc_type is the exception's type (its class).
- exc_value is the exception value itself. It's the same value that we'd get when catching the exception via except SomeErrorClass as exc_value:.
- traceback is a traceback object, including stack information. This has the same information that we see when a traceback is printed, like line numbers and source file names.
These are useful when we need to make decisions based on the exact exception that happened. Or we might write a context manager whose job is to log exceptions and their tracebacks.
Normally, the exception "escapes" from the with block, bubbles up, and becomes an error. For example, when using files as context managers, exceptions will always escape.

from io import StringIO

with StringIO() as file:
  raise ValueError("it blew up")

Result:

ValueError: it blew up

If .__exit__ returns True, that silences the exception, preventing it from escaping the with.
The next example defines a context manager that ignores ValueErrors, but doesn't ignore other exceptions.

class IgnoreValueError:
  def __enter__(self):
    # This method does nothing, but it must be here to use instances as context
    # managers.
    pass

  def __exit__(self, exc_type, exc_value, traceback):
    if isinstance(exc_value, ValueError):
      return True
    else:
      return False

Our context manager will silence ValueError (by returning True), but let other exceptions escape. As usual, a raise stops execution at that point, so the rest of the with block doesn't run.

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

breadcrumb = 1
with IgnoreValueError():
  breadcrumb = 2
  raise ValueError()
  breadcrumb = 3
breadcrumb

Result:

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

with IgnoreValueError():
  raise TypeError("this is a type error, not a value error")

Result:

TypeError: this is a type error, not a value error

A surprising amount of programming takes the form of "set up, do something, then clean up." That makes the with statement quite versatile. Many use cases are examples of resource management, like managing files or network connections. But we can use with in any situation where we have a block of code, and we need some other code to run before or after it.