Decorators
A decorator wraps a function to add behaviour before or after it runs, without changing the function's own code. You've already used one: @dataclass. Others you'll see constantly include @property, @app.route(...) in Flask, and @lru_cache.
Functions are objects
Decorators rely on the fact that functions can be passed around and created inside other functions:
def shout(text):
return text.upper()
def apply(func, value):
return func(value)
print(apply(shout, "hello"))A decorator by hand
A decorator takes a function and returns a new function that wraps it:
def announce(func):
def wrapper():
print(f"About to run {func.__name__}")
func()
print(f"Finished {func.__name__}")
return wrapper
def greet():
print("Hello!")
greet = announce(greet) # replace greet with the wrapped version
greet()The @ syntax
@announce above a definition is shorthand for greet = announce(greet):
def announce(func):
def wrapper():
print(f"About to run {func.__name__}")
func()
print(f"Finished {func.__name__}")
return wrapper
@announce
def greet():
print("Hello!")
greet()Wrapping functions that take arguments
Use *args, **kwargs so the wrapper accepts and passes on anything, and return the result:
import functools
import time
def timer(func):
@functools.wraps(func) # keeps the original name and docstring
def wrapper(*args, **kwargs):
start = time.perf_counter()
result = func(*args, **kwargs)
ms = (time.perf_counter() - start) * 1000
print(f"{func.__name__} took {ms:.2f} ms")
return result
return wrapper
@timer
def total_squares(n):
return sum(i * i for i in range(n))
print(total_squares(200_000))
print(total_squares.__name__)A practical one: caching
functools.lru_cache remembers results, so repeated calls with the same arguments are instant. It turns the slow recursive Fibonacci into a fast one:
from functools import lru_cache
@lru_cache(maxsize=None)
def fib(n):
return n if n < 2 else fib(n - 1) + fib(n - 2)
print(fib(80))
print(fib.cache_info())Without the cache, fib(80) would take longer than the age of the universe; with it, it takes microseconds.
Decorators with arguments
A decorator that takes settings needs one more layer: a function that returns the decorator:
def repeat(times):
def decorator(func):
def wrapper(*args, **kwargs):
for _ in range(times):
func(*args, **kwargs)
return wrapper
return decorator
@repeat(3)
def cheer(name):
print(f"Go {name}!")
cheer("Asha")