Writing Custom Decorators

Introduction

What is a Decorator?

def shout(func):
    def wrapper(*a, **kw):
        result = func(*a, **kw)
        return str(result).upper()
    return wrapper
@shout
def greet(name): return f'hello {name}'
print(greet('Alice'))

The Wrapper Pattern

def timer(func):
    import time
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        print(f'{func.__name__} took {time.time()-start:.4f}s')
        return result
    return wrapper
@timer
def slow(): import time; time.sleep(0.01)
slow()

functools.wraps

from functools import wraps
def log(func):
    @wraps(func)
    def wrapper(*a, **kw):
        print(f'calling {func.__name__}')
        return func(*a, **kw)
    return wrapper
@log
def add(a, b): return a + b
print(add.__name__)
add(1, 2)

Decorator with Arguments

def repeat(n):
    def decorator(func):
        from functools import wraps
        @wraps(func)
        def wrapper(*a, **kw):
            for _ in range(n): func(*a, **kw)
        return wrapper
    return decorator
@repeat(3)
def hi(): print('hi')
hi()

Stacking Decorators

def bold(f):
    from functools import wraps
    @wraps(f)
    def w(*a,**k): return '<b>' + f(*a,**k) + '</b>'
    return w
def italic(f):
    from functools import wraps
    @wraps(f)
    def w(*a,**k): return '<i>' + f(*a,**k) + '</i>'
    return w
@bold
@italic
def text(): return 'hello'
print(text())

Class-Based Decorator

class Counter:
    def __init__(self, func): self.func,self.count=func,0
    def __call__(self, *a, **kw):
        self.count+=1
        return self.func(*a, **kw)
@Counter
def greet(name): return f'hi {name}'
greet('Alice'); greet('Bob')
print(greet.count)

Memoization Decorator

from functools import lru_cache
@lru_cache(maxsize=None)
def fib(n):
    if n < 2: return n
    return fib(n-1) + fib(n-2)
print(fib(30))

Authorization Decorator

def requires_auth(func):
    from functools import wraps
    @wraps(func)
    def wrapper(user, *a, **kw):
        if not user.get('authenticated'):
            raise PermissionError('Not authenticated')
        return func(user, *a, **kw)
    return wrapper
@requires_auth
def dashboard(user): return 'Welcome!'
try: dashboard({'authenticated': False})
except PermissionError as e: print(e)

Retry Decorator

import random
from functools import wraps
def retry(times=3):
    def dec(func):
        @wraps(func)
        def wrapper(*a, **kw):
            for i in range(times):
                try: return func(*a, **kw)
                except Exception as e:
                    if i==times-1: raise
        return wrapper
    return dec
@retry(3)
def flaky():
    if random.random()<0.7: raise ValueError
    return 'ok'
try: print(flaky())
except ValueError: print('failed')

Decorator Best Practices

from functools import wraps
def validate_positive(func):
    @wraps(func)
    def wrapper(x, *a, **kw):
        if x <= 0: raise ValueError(f'{x} must be positive')
        return func(x, *a, **kw)
    return wrapper
@validate_positive
def sqrt(x): return x**0.5
print(sqrt(9))

Quick Check

Recap

Keep Going