Lesson 13 · 13 of 20

Custom exceptions and context managers · Python

Create your own exceptions and context managers in Python: with, __enter__, __exit__, contextlib.contextmanager and suppress, in a warehouse.

Python · Custom exceptions and context managers in DMarlon Studio: the example's code and its output in the terminal

A custom exception is a class that inherits from Exception: class StockError(Exception): pass is enough. It lets you catch exactly that problem (running out of stock) without mixing it up with other errors. take() raises it with a message saying what was missing.

A context manager is whatever goes after with. Transaction is one because it has two methods: `__enter__` runs on the way in (it saves a copy of the stock and returns the dictionary, which reaches s through as) and `__exit__` runs on the way out, error or not.

If an exception is raised inside the block, __exit__ receives it in exc_type and exc. Here it restores the copy with update(), so the 3 pens already taken go back: either the whole order goes through or none of it does. Returning True tells Python the error has been dealt with; that is why only StockError is silenced and any other error would carry on up.

`@contextmanager` builds a manager from a generator: the code before yield is the way in and the code after it is the way out. section() prints the opening and closing tags. A single with can open several managers separated by commas. They close in reverse order on the way out, first Transaction and then section: that is why the closing tag comes after the rollback message.

`suppress(KeyError)` is a ready-made manager that ignores one kind of error: asking for "paper", which is not in stock, does not stop the program.

Code · main.py

from contextlib import contextmanager, suppressclass StockError(Exception): pass     # our own kind of errorclass Transaction:                    # a context manager written as a class    def __init__(self, stock):        self.stock = stock    def __enter__(self):        self.backup = dict(self.stock)        return self.stock    def __exit__(self, exc_type, exc, tb):        if exc_type is not None:      # something failed: undo everything            self.stock.update(self.backup)            print(f"  rolled back: {exc}")        return exc_type is StockError     # True = error handled here@contextmanager                       # the same idea with a generatordef section(title):    print(f"[{title}]")    yield    print(f"[/{title}]")def take(stock, product, qty):    if stock[product] < qty:        raise StockError(f"{product}: wanted {qty}, only {stock[product]} left")    stock[product] -= qtystock = {"pen": 10, "ink": 2}with section("order 1"), Transaction(stock) as s:    take(s, "pen", 3)    take(s, "ink", 5)                 # fails: the pens come back toowith suppress(KeyError):              # ignore one specific error    take(stock, "paper", 1)print("Stock:", stock)

Output · python3 main.py

[order 1]
  rolled back: ink: wanted 5, only 2 left
[/order 1]
Stock: {'pen': 10, 'ink': 2}

Key points

A custom exception inherits from Exception.
with calls __enter__ on entry and __exit__ on exit, always.
If __exit__ returns True, the exception does not propagate.
@contextmanager builds a manager from a generator with yield.
contextlib.suppress ignores the errors you name.

Exercise

Place a second order that can be fulfilled (2 pens and 1 ink) inside section("order 2") and check that the stock changes this time.