Writing Context Managers with __enter__ and __exit__
22 min · Intermediate
Need help? Ask a tutor on WhatsApp (opens in a new tab)Any class can work with with. It needs two methods: one for setup and one for cleanup. This is the tool for your own resources, such as a transaction or a temporary setting.
Two methods make a context manager
Context manager protocol is the pair of methods __enter__ and __exit__ that a class defines so that with can set it up and clean it up.
__enter__ runs when the block starts. __exit__ runs when it ends. Define both on a class and with accepts it.
class Section: def __enter__(self): print("enter") return self
def __exit__(self, exc_type, exc, tb): print("exit")
with Section(): print("body")Example 1 · Order of enter, body, exit
class Section: def __init__(self, name): self.name = name
def __enter__(self): print("enter", self.name) return self
def __exit__(self, exc_type, exc, tb): print("exit", self.name)
with Section("outer"): with Section("inner"): print("body")enter outer enter inner body exit inner exit outer
What this code does
- Nested managers exit in reverse order: inner first.
- Each section prints
exiteven though nothing failed.
The output is enter, body, exit. The three arguments of __exit__ are used in a later section.
What __enter__ returns
The value after as is the return value of __enter__. Return self when the manager is also the resource, or return another object, such as a connection.
class Connection: def __enter__(self): self.open = True return self
def __exit__(self, exc_type, exc, tb): self.open = False
with Connection() as conn: print(conn.open)print(conn.open)Example 2 · The as target
class Basket: def __enter__(self): self.items = [] return self.items
def __exit__(self, exc_type, exc, tb): print("closing basket")
with Basket() as items: items.append("pen") items.append("book") print(items)['pen', 'book'] closing basket
What this code does
__enter__returns the list, soitemsis the list, not the manager.- Returning a different object than
selfis common.
The object is open inside the block and closed after it. If __enter__ returns nothing, as would receive None.
__exit__ receives the exception
Exception suppression means __exit__ returns True, which tells Python to treat the exception as handled so the code after the with block continues normally.
If the block raised an error, __exit__ gets its type, value and traceback. If not, all three are None. Return True to hide the error. Return False or nothing to let it continue.
class Audit: def __enter__(self): return self
def __exit__(self, exc_type, exc, tb): if exc_type: print("failed:", exc_type.__name__, exc) return FalseExample 3 · Reading the exception in __exit__
class Audit: def __enter__(self): return self
def __exit__(self, exc_type, exc, tb): if exc_type: print("failed:", exc_type.__name__, exc) else: print("ok") return False
with Audit(): pass
try: with Audit(): int("abc")except ValueError: print("caller saw the error")ok failed: ValueError invalid literal for int() with base 10: 'abc' caller saw the error
What this code does
- On success, the three arguments are
None, so theelsebranch printsok. - Returning
Falselets theValueErrorreach theexceptline.
Example 4 · Suppressing one exception type
class Ignore: def __init__(self, *errors): self.errors = errors
def __exit__(self, exc_type, exc, tb): return exc_type is not None and issubclass(exc_type, self.errors)
def __enter__(self): return self
with Ignore(KeyError): {}["missing"]print("continued")
try: with Ignore(KeyError): 1 / 0except ZeroDivisionError: print("not hidden")continued not hidden
What this code does
- The
KeyErroris hidden because__exit__returnedTrue. - A
ZeroDivisionErroris not in the list, so__exit__returnsFalseand it escapes.
Returning False is the safe default: the error is logged here and then still reaches the caller.
A transaction manager
A transaction should save its changes only if everything worked. A context manager fits: commit when the block ends cleanly, roll back when it raises.
class Transaction: def __init__(self, db): self.db = db
def __enter__(self): self.start = list(self.db) return self.db
def __exit__(self, exc_type, exc, tb): if exc_type: self.db[:] = self.startExample 5 · Commit or roll back
class Transaction: def __init__(self, db): self.db = db
def __enter__(self): self.start = list(self.db) return self.db
def __exit__(self, exc_type, exc, tb): if exc_type: self.db[:] = self.start
accounts = ["A:100"]with Transaction(accounts) as db: db.append("B:50")print(accounts)
try: with Transaction(accounts) as db: db.append("C:75") raise RuntimeError("payment failed")except RuntimeError: passprint(accounts)['A:100', 'B:50'] ['A:100', 'B:50']
What this code does
- The first block ends cleanly, so the change stays.
- The second raises an error, so
C:75is removed.
On an error, the list is put back to the state saved in __enter__, and the error still continues because __exit__ returns None.
Next step
Check your understanding
12 quick questions