lesson 161 Decorators & Context Managers · advanced

Decorators: wrap a function

A decorator takes a function and returns a replacement, usually a wrapper that runs some code around the original: counting calls, logging, checking arguments. The @ line above a def is shorthand for passing the function through it.

You'll be able to

Explain that @deco above def f means f = deco(f)
Write a wrapper that accepts any arguments, calls the original and returns its result
Keep state for a decorated function and know what happens when the original raises

The idea

Functions are values (see Advanced Functions & Recursion): you can pass one to another function, and a function can build and return a new one that remembers its surroundings (a closure). A decorator uses exactly that. It takes a function, defines a wrapper that does something extra and calls the original, and returns the wrapper. The wrapper accepts *args (a tuple of positional arguments) and **kwargs (a dictionary of keyword arguments) and passes them on unchanged, so it fits any function.

Writing @count_calls on the line above def parse_magnitude(...) is shorthand for parse_magnitude = count_calls(parse_magnitude): the name now refers to the wrapper.

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

example-1.py · Counting calls, by hand and with @

def count_calls(func):
    def wrapper(*args, **kwargs):
        wrapper.calls += 1
        return func(*args, **kwargs)
    wrapper.calls = 0
    return wrapper

def parse_magnitude(text):
    return float(text)

parse_magnitude = count_calls(parse_magnitude)
print(parse_magnitude("5.5"), parse_magnitude(" 4.7 "))
print(parse_magnitude.calls)

@count_calls
def place_name(feature):
    return feature["properties"]["place"]

print(place_name({"properties": {"place": "South Sandwich Islands region"}}))
print(place_name.calls)

A wrapper runs code before and after the call, and it sees what came back, or the exception that came out. This logging decorator records each call of a parser, including the failures, and then re-raises the error so the caller still knows something went wrong.

example-2.py · Log every call, including failures

log = []

def log_calls(func):
    def wrapper(*args, **kwargs):
        try:
            result = func(*args, **kwargs)
        except Exception as err:
            log.append(f"{func.__name__}{args} raised {type(err).__name__}")
            raise
        log.append(f"{func.__name__}{args} -> {result!r}")
        return result
    return wrapper

@log_calls
def parse_depth(text, unit="km"):
    value = float(text)
    return value / 1000 if unit == "m" else value

parse_depth("22")
parse_depth("10", unit="m")
try:
    parse_depth("n/a")
except ValueError:
    pass
for line in log:
    print(line)

Forget the last line of a decorator and the decorated name becomes None:

example-3.py · A decorator that returns nothing

def count_calls(func):
    def wrapper(*args, **kwargs):
        wrapper.calls += 1
        return func(*args, **kwargs)
    wrapper.calls = 0

@count_calls
def parse_magnitude(text):
    return float(text)

print(parse_magnitude("5.5"))

When to use it: behaviour that many functions need in the same way and that isn't their main job: counting, logging, timing, caching, checking arguments. Mistakes to avoid: forgetting return wrapper; a wrapper that calls the function but doesn't return its result; a wrapper with fixed parameters instead of *args, **kwargs; and swallowing the original's exceptions.

Review

A decorator takes a function and returns a replacement; @deco above def f means f = deco(f).
A wrapper takes *args, **kwargs, calls the original with them, and returns its result.
State such as a count or a cache can live on the wrapper; let exceptions through unless handling them is the decorator's job.
Next: decorators that take their own settings, like @retry_on(TimeoutError, attempts=3).