Context Managers Checkpoint

25 min · Intermediate

Fix the mistakes

The slips that cost marks in this module. Spot each one, then check the fix.

A real export job opens files, changes a setting for a moment and must clean up after any failure. This checkpoint builds that job from context managers and tests the whole module.

Pick the simplest manager

Use a built-in manager when one exists, such as open or a lock. Use @contextmanager for short single-use setup and cleanup. Use a class when you need attributes, reuse or fine control of errors.

Python

1from contextlib import contextmanager2
3@contextmanager4def export_mode(settings):5    settings["exporting"] = True6    try:7        yield settings8    finally:9        settings["exporting"] = False

Example 1 · Export mode with rollback on error

Python

1from contextlib import contextmanager2
3settings = {"exporting": False}4rows = []5
6@contextmanager7def export_mode():8    settings["exporting"] = True9    saved = list(rows)10    try:11        yield rows12    except Exception:13        rows[:] = saved14        raise15    finally:16        settings["exporting"] = False17
18with export_mode() as r:19    r.append("order 1")20print(rows, settings["exporting"])21
22try:23    with export_mode() as r:24        r.append("order 2")25        raise RuntimeError("write failed")26except RuntimeError:27    print(rows, settings["exporting"])

Output

['order 1'] False
['order 1'] False

What this code does

finally always switches the setting off; except rolls the rows back and re-raises.
The caller still sees the RuntimeError.

The setting is turned on for the block and always switched off, even after a failure.

Nest managers and read the order

Managers enter from left to right and exit from right to left. The first resource opened is the last one closed, so later ones can safely rely on earlier ones.

Python

with section("outer"):    with section("inner"):        print("body")

Example 2 · Order of nested managers

Python

1from contextlib import contextmanager2
3@contextmanager4def section(name):5    print("enter", name)6    try:7        yield8    finally:9        print("exit", name)10
11with section("outer"), section("inner"):12    print("body")

Output

enter outer
enter inner
body
exit inner
exit outer

What this code does

Two managers in one with line behave like nested blocks.
The cleanup order is the reverse of the setup order.

The order is: enter outer, enter inner, body, exit inner, exit outer.

Decide about errors

Let errors continue by default. Suppress only a named, expected error with suppress(...), and keep its block small. A manager that rolls back should still re-raise.

Python

with suppress(FileNotFoundError):    os.remove("old-report.csv")

Example 3 · A class manager with state

Python

1class Counter:2    opened = 03
4    def __enter__(self):5        Counter.opened += 16        return self7
8    def __exit__(self, exc_type, exc, tb):9        return exc_type is KeyError10
11with Counter():12    {}["x"]13print("still running", Counter.opened)14try:15    with Counter():16        1 / 017except ZeroDivisionError:18    print("passed through", Counter.opened)

Output

still running 1
passed through 2

What this code does

The KeyError is suppressed because the check returns True.
A ZeroDivisionError returns False, so it continues to the caller.

If the file is missing, the program carries on. Any other error still stops the program.

Finish up

Mark complete & continue

Next: Annotating Functions and Classes

1from contextlib import contextmanager2 3settings = {"exporting": False}4rows = []5 6@contextmanager7def export_mode():8 settings["exporting"] = True9 saved = list(rows)10 try:11 yield rows12 except Exception:13 rows[:] = saved14 raise15 finally:16 settings["exporting"] = False17 18with export_mode() as r:19 r.append("order 1")20print(rows, settings["exporting"])21 22try:23 with export_mode() as r:24 r.append("order 2")25 raise RuntimeError("write failed")26except RuntimeError:27 print(rows, settings["exporting"])

1from contextlib import contextmanager2 3@contextmanager4def section(name):5 print("enter", name)6 try:7 yield8 finally:9 print("exit", name)10 11with section("outer"), section("inner"):12 print("body")

1class Counter:2 opened = 03 4 def __enter__(self):5 Counter.opened += 16 return self7 8 def __exit__(self, exc_type, exc, tb):9 return exc_type is KeyError10 11with Counter():12 {}["x"]13print("still running", Counter.opened)14try:15 with Counter():16 1 / 017except ZeroDivisionError:18 print("passed through", Counter.opened)