Writing Your Own Context Manager
In short: a context manager is any object with __enter__ and
__exit__; with mgr as v: calls __enter__ and binds its return value to
v, then calls __exit__ no matter how the block ends — on a clean
exit its three arguments are None, on an exception they describe it, and
if __exit__ returns something truthy the exception is swallowed.
contextlib.contextmanager turns a generator with one yield into the same
thing: everything before yield is __enter__, the yielded value is what
as binds, everything after is __exit__ — but it must sit in a
try/finally, because an exception from the body is thrown back in at the
yield, and code after it only runs if a finally catches that detour.
Several managers exit in reverse order, last-entered first, same as nested
with blocks.
Announce a stage, pass or fail
A pipeline stage should report whether it succeeded, without every stage
writing that logic by hand. A class with __enter__ and __exit__ does it
once:
The parse stage reported its failure and still let the exception
through, because __exit__ returned None. That's the right default — a
manager cleans up and reports; it doesn't decide the error doesn't matter.
Returning something truthy is an opt-in to silence, and it's easy to get
wrong by accident:
class Hush:
def __enter__(self):
return self
def __exit__(self, exc_type, exc, tb):
return True
with Hush():
1 / 0
print("inside")
print("after")Open every file in a batch at once
Quarry sometimes needs to open every file in a batch together — and the
batch size comes from the data, so with a, b, c: won't do.
contextlib.ExitStack holds a variable number of managers and unwinds them
all when its own with ends:
enter_context enters a manager now and schedules its exit for later. If
opening the third file had failed partway through, the first two would
still be closed — nothing leaks. When several managers are known ahead of
time, plain nesting unwinds the same way, last-in first-out:
log = []
class Tag:
def __init__(self, name):
self.name = name
def __enter__(self):
log.append("enter " + self.name)
def __exit__(self, exc_type, exc, tb):
log.append("exit " + self.name)
with Tag("a"), Tag("b"):
log.append("body")
print(", ".join(log))Here is the same unwinding as a diagram — push on the way in, pop on the way out, and an exception takes the same exit route:
Rewrite a report without losing the old one on a crash
Last week Quarry's machine lost power while rewriting its summary report.
The file was opened with "w" — which empties it immediately — and only
half the new report reached disk. The fix: never write over the real file.
Write into a temporary sibling, and only when it's complete, rename it
over the original with os.replace, which is atomic — any reader sees
either the whole old file or the whole new one, never a half-written one.
Writing a class with two dunder methods for a five-line cleanup is a lot of
ceremony. @contextmanager turns a generator into the same thing. A
generator is a function that can pause — calling it doesn't run it; each
resume runs up to the next yield and parks there with its locals intact:
That try/finally isn't decoration. If the with body raises,
@contextmanager throws that exception into the generator, at the
yield line. Without a finally, the lines after yield never run, and
the cleanup silently goes missing:
from contextlib import contextmanager
log = []
@contextmanager
def opened(name):
log.append("open " + name)
yield name
log.append("close " + name)
try:
with opened("a.log"):
log.append("body")
raise ValueError("bad row")
except ValueError:
pass
print(log)Quarry's guard follows the same shape: before the yield, open the .tmp
sibling; yield the open file for the caller to write into; if the body
finished, close it and os.replace it over the original; if it raised,
close it, delete the .tmp, and re-raise — the guard cleans up, it
doesn't pretend the write succeeded.
Quick practice
Run it. Only [4] prints, not [4, 10]. Move the with so each value gets its own chance.
from contextlib import closing
class Handle:
def __init__(self):
self.closed = False
def close(self):
self.closed = True
h = Handle()
with closing(h) as handle:
pass
print(handle is h, h.closed)class Counter:
def __enter__(self):
return 0
def __exit__(self, exc_type, exc, tb):
return False
with Counter() as n:
n += 1
print(n)from contextlib import contextmanager
@contextmanager
def guarded():
try:
yield "ok"
except ZeroDivisionError:
yield "recovered"
with guarded() as v:
1 / 0
print(v)🎯 Your turn
Write replacing(path):
atomic_write("old", "new", False)→{"content": "new", "files": ["report.txt"]}atomic_write("old report", "new report", True)→{"content": "old report", "files": ["report.txt"]}
Hint — tmp = path.with_name(path.name + ".tmp"); wrap yield f in
try/except BaseException/else; success replaces, failure unlinks and re-raises.
Then press ▶ Run and hit ✓ Check. ✅