Context Managers 🚪

The problem it solves

# Setup, work, cleanup. The cleanup must happen even if the work explodes.
print("open the door")
try:
    print("do the work")
    raise ValueError("something went wrong")
except ValueError as err:
    print("caught:", err)
finally:
    print("close the door")

open the door
do the work
caught: something went wrong
close the door

That try/finally is correct and tedious, and you have to remember it every single time. A context manager packages the pattern so the caller cannot forget.

Writing one with a class

class Door:
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        print(f"opening {self.name}")
        return self                # whatever `as` binds to

    def __exit__(self, exc_type, exc_value, traceback):
        print(f"closing {self.name}")
        return False               # False means "do not swallow the exception"

    def knock(self):
        return "nobody answers"


with Door("the vault") as door:
    print(door.knock())

print("---")

try:
    with Door("the trapdoor") as door:
        raise RuntimeError("floor gives way")
except RuntimeError as err:
    print("caught outside:", err)

opening the vault
nobody answers
closing the vault
---
opening the trapdoor
closing the trapdoor
caught outside: floor gives way

Note the second case: the exception fired inside the block, __exit__ still ran, and only then did the error continue outward. That guarantee is the entire point.

The easy way: @contextmanager

from contextlib import contextmanager


@contextmanager
def door(name):
    print(f"opening {name}")
    try:
        yield name              # everything before this is __enter__
    finally:
        print(f"closing {name}")   # everything after is __exit__


with door("the hatch") as which:
    print(f"inside {which}")

opening the hatch
inside the hatch
closing the hatch

One generator, one yield. Setup above, cleanup below, wrapped in try/finally so the cleanup survives an exception. This is how most context managers are written in practice.

A genuinely useful one: timing a block

import time
from contextlib import contextmanager


@contextmanager
def timer(label):
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start
        print(f"{label}: finished in under a second: {elapsed < 1}")


with timer("summing a million"):
    total = sum(range(1_000_000))

print(f"{total:,}")

summing a million: finished in under a second: True
499,999,500,000

Temporarily changing something, and putting it back

import os
from contextlib import contextmanager


@contextmanager
def env(**changes):
    """Set environment variables for the duration of the block."""
    original = {k: os.environ.get(k) for k in changes}
    os.environ.update({k: str(v) for k, v in changes.items()})
    try:
        yield
    finally:
        for key, value in original.items():
            if value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = value


print("before:", os.environ.get("SHIP", "not set"))

with env(SHIP="Sea Monkey"):
    print("inside:", os.environ["SHIP"])

print("after: ", os.environ.get("SHIP", "not set"))

before: not set
inside: Sea Monkey
after:  not set

contextlib's ready-made tools

from contextlib import suppress, redirect_stdout
import io
from pathlib import Path

# 1. suppress: a try/except/pass that is honest about being one
with suppress(FileNotFoundError):
    Path("nope.txt").unlink()
print("no explosion")

# 2. redirect_stdout: capture prints
buffer = io.StringIO()
with redirect_stdout(buffer):
    print("this goes into the buffer")
print("captured:", buffer.getvalue().strip())

no explosion
captured: this goes into the buffer

Several at once

from pathlib import Path

Path("in.txt").write_text("grog\nmap\nsword\n", encoding="utf-8")

with (
    open("in.txt", encoding="utf-8") as source,
    open("out.txt", "w", encoding="utf-8") as target,
):
    for line in source:
        target.write(line.upper())

print(Path("out.txt").read_text(encoding="utf-8").strip())

GROG
MAP
SWORD

Both files are guaranteed closed, in reverse order, whatever happens. The bracketed multi-line form needs Python 3.10 or newer; before that you separated them with commas on one long line.

The catch: exceptions during cleanup

from contextlib import contextmanager


@contextmanager
def careless():
    yield
    print("this cleanup NEVER runs when the block raises")


@contextmanager
def careful():
    try:
        yield
    finally:
        print("this cleanup always runs")


for manager in (careless, careful):
    try:
        with manager():
            raise ValueError("boom")
    except ValueError:
        print(f"  ({manager.__name__} finished)")

  (careless finished)
this cleanup always runs
  (careful finished)

If you write @contextmanager without try/finally, the code after yield is skipped whenever the block raises, which is precisely the case you were trying to protect. The try/finally is not optional decoration.

import os
from contextlib import contextmanager
from pathlib import Path


@contextmanager
def working_directory(path):
    """Change directory for the duration of the block, then change back."""
    previous = Path.cwd()
    Path(path).mkdir(parents=True, exist_ok=True)
    os.chdir(path)
    try:
        yield Path(path)
    finally:
        os.chdir(previous)


start = Path.cwd().name

with working_directory("cargo/hold"):
    print("inside:", Path.cwd().name)
    Path("manifest.txt").write_text("47 barrels\n", encoding="utf-8")

print("back where we started:", Path.cwd().name == start)
print(Path("cargo/hold/manifest.txt").read_text(encoding="utf-8").strip())

inside: hold
back where we started: True
47 barrels

from contextlib import contextmanager
import copy


@contextmanager
def transaction(data):
    """Work on a copy; only commit it if the block finishes cleanly."""
    working = copy.deepcopy(data)
    yield working
    data.clear()
    data.update(working)


accounts = {"Guybrush": 100, "Elaine": 250}

with transaction(accounts) as draft:
    draft["Guybrush"] -= 50
    draft["Elaine"] += 50
print("committed:", accounts)

try:
    with transaction(accounts) as draft:
        draft["Guybrush"] -= 500
        raise ValueError("insufficient funds")
except ValueError as err:
    print("rolled back:", err)

print("unchanged: ", accounts)

committed: {'Guybrush': 50, 'Elaine': 300}
rolled back: insufficient funds
unchanged:  {'Guybrush': 50, 'Elaine': 300}

# Their version
def read_config(path):
    f = open(path, encoding="utf-8")
    data = f.read()
    value = int(data)          # if this raises ValueError...
    f.close()                  # ...this line never runs
    return value


# The point: the close is skipped on any exception, and on any early return.
try:
    read_config("/dev/null")
except ValueError:
    print("the file object is now leaked until garbage collection")

# The version that cannot leak
def read_config_safely(path):
    with open(path, encoding="utf-8") as f:
        data = f.read()
    return int(data)


try:
    read_config_safely("/dev/null")
except ValueError:
    print("file already closed, guaranteed, before this line ran")

the file object is now leaked until garbage collection
file already closed, guaranteed, before this line ran