lesson 164 Decorators & Context Managers · advanced

Context managers: enter and exit

A with block guarantees that cleanup runs: a file closes, a transaction ends, a setting is restored, even when the block raises. A class becomes a context manager by defining __enter__, which runs at the start, and __exit__, which always runs at the end and sees any exception.

You'll be able to

Trace a with statement: __enter__, the block, then __exit__(exc_type, exc, traceback)
Restore state in __exit__ whether or not the block failed
Suppress only chosen exceptions by returning True from __exit__

The idea

You've used with open(...) (see Files & Directories) and with conn: (see Packaging & Databases). Both work through the context-manager protocol. with manager as name: calls manager.__enter__() and binds what it returns to name. Then the block runs. Then manager.__exit__(exc_type, exc, traceback) runs no matter how the block ended: normally (all three arguments are None), with an exception (they describe it), or through return or break. If __exit__ returns True, an exception from the block is suppressed; otherwise it continues after cleanup.

Run each example to see what it prints. Running examples doesn't count toward progress.

example-1.py · Change a setting for one block

settings = {"currency": "USD", "timeout": 10}

class TempSetting:
    def __init__(self, settings, key, value):
        self.settings = settings
        self.key = key
        self.value = value

    def __enter__(self):
        self.old = self.settings[self.key]
        self.settings[self.key] = self.value
        return self.settings

    def __exit__(self, exc_type, exc, traceback):
        self.settings[self.key] = self.old
        return False

with TempSetting(settings, "currency", "EUR") as current:
    print("inside:", current["currency"])
print("after:", settings["currency"])

try:
    with TempSetting(settings, "timeout", 1):
        raise TimeoutError("the export took too long")
except TimeoutError as err:
    print("error:", err)
print("timeout restored:", settings["timeout"])

__exit__ receives the exception, so it can decide what happens to it. This collector suppresses only the error types it was given and keeps their messages, which lets an import continue past bad rows while every other kind of error still stops it.

example-2.py · Collect chosen errors and carry on

class Collect:
    def __init__(self, *types):
        self.types = types
        self.errors = []

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, traceback):
        if exc_type is not None and issubclass(exc_type, self.types):
            self.errors.append(f"{exc_type.__name__}: {exc}")
            return True
        return False

rows = ["4.7", "n/a", "5.5", ""]
magnitudes = []
collector = Collect(ValueError)
for row in rows:
    with collector:
        magnitudes.append(float(row))
print(magnitudes)
print(collector.errors)

When to use it: anything that must be undone or closed: temporary settings, locks, working directories, timers, staged changes. Mistakes to avoid: doing the undo outside __exit__ (it's skipped when the block raises); returning True from __exit__ by accident, which hides every error; forgetting to return self (or the resource) from __enter__ when the block uses as; and setting up in __init__ work that belongs in __enter__.

Review

with m as x: calls m.__enter__() (its result is x), runs the block, then always calls m.__exit__.
Undo work in __exit__, which receives the exception type, value and traceback (all None on success).
Return True from __exit__ only to suppress exceptions the manager is meant to handle.
Next: contextlib, which writes context managers as generators and adds ready-made ones.

Context managers: __enter__ and __exit__

Context managers: enter and exit