CodeMIRAI
0 XP

Writing context managers

2 min read 5 runnable examples Quiz included

You've used with open(...) to make sure files close. Any object that supports with is a context manager, and you can write your own for anything that needs setting up and tidying away.

As a class: __enter__ and __exit__

Python
class Timer:
    def __enter__(self):
        import time
        self.start = time.perf_counter()
        return self                    # this is what "as t" receives

    def __exit__(self, exc_type, exc_value, traceback):
        import time
        self.ms = (time.perf_counter() - self.start) * 1000
        print(f"Block took {self.ms:.2f} ms")
        return False                   # don't hide exceptions

with Timer() as t:
    total = sum(i * i for i in range(300_000))
print(total)
  • __enter__ runs at the start of the with block; its return value goes to as.
  • __exit__ runs at the end, even if the block raised an exception. It receives the exception details (or Nones).

Exceptions inside the block

Python
class Section:
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        print(f"[enter {self.name}]")

    def __exit__(self, exc_type, exc, tb):
        status = "with an error" if exc_type else "cleanly"
        print(f"[exit {self.name} {status}]")
        return False

try:
    with Section("upload"):
        print("uploading...")
        raise ConnectionError("network dropped")
except ConnectionError as e:
    print("Caught:", e)

If __exit__ returned True, the exception would be swallowed. Only do that deliberately.

The easy way: @contextmanager

contextlib.contextmanager turns a generator into a context manager. Code before yield is the setup, code after it is the tidy-up:

Python
from contextlib import contextmanager

@contextmanager
def tag(name):
    print(f"<{name}>")
    yield
    print(f"</{name}>")

with tag("ul"):
    with tag("li"):
        print("Item one")

Wrap the yield in try/finally so the tidy-up always runs:

Python
from contextlib import contextmanager
import os

@contextmanager
def temporary_file(name, text):
    with open(name, "w") as f:
        f.write(text)
    try:
        yield name
    finally:
        os.remove(name)
        print(f"(removed {name})")

with temporary_file("draft.txt", "hello") as path:
    print(open(path).read())
print(os.path.exists("draft.txt"))

Handy ones in the standard library

Python
from contextlib import suppress, redirect_stdout
import io

with suppress(FileNotFoundError):          # ignore one specific error
    open("missing.txt")
print("carried on")

buffer = io.StringIO()
with redirect_stdout(buffer):              # capture print output
    print("captured!")
print("buffer holds:", buffer.getvalue().strip())
Finished this lesson?Mark it complete to earn +20 XP and move your progress on.