ruk·si

🐍 Python
Decorators

Updated at 2018-06-09 15:56

Decorators are commonly used in:

  • logging
  • access control
  • timing
  • rate-limiting
  • caching

Decorators wrap a function around another function. Use @functools.wraps decorator on the wrapper function to pass through the metadata of the underlying function.

from functools import wraps
from typing import Callable, TypeVar

F = TypeVar('F', bound=Callable[[], str])

def strong(func: F) -> F:
    @wraps(func)
    def wrapper() -> str:
        return '<strong>' + func() + '</strong>'

    return wrapper

def emphasis(func: F) -> F:
    @wraps(func)
    def wrapper() -> str:
        return '<em>' + func() + '</em>'

    return wrapper

@strong
@emphasis
def greet() -> str:
    """ Returns a friendly greeting. """
    return 'Hello!'

assert greet() == '<strong><em>Hello!</em></strong>'
assert greet.__name__ == 'greet'
assert greet.__doc__ == ' Returns a friendly greeting. '

Remember to pass through all the parameters of the wrapped function.

from functools import wraps
from typing import Callable, TypeVar

F = TypeVar('F', bound=Callable)

def trace(func: F) -> F:
    @wraps(func)
    def wrapper(*args, **kwargs):
        print(f'TRACE: calling {func.__name__}() with {args}, {kwargs}')
        original_result = func(*args, **kwargs)
        return original_result

    return wrapper

@trace
def say(name: str, line: str, extra: str) -> str:
    """ Returns a piece of dialogue. """
    return f'{name}: {line}'

print(say('Jane', 'Hello!', extra='whisper'), say.__name__, say.__doc__)
# => TRACE: calling say() with ('Jane', 'Hello!'), {'extra': 'whisper'}
# => Jane: Hello! say  Returns a piece of dialogue.

Decorators are called right after the decorator function is bound. The wrapped function doesn't even have to be called for the decorator to run. Decorator functions run even if the file is just imported.

from typing import Callable, List

registry: List[Callable] = []

def register(func: Callable) -> Callable:
    # we don't need functools.wrap here as function is returned as is
    registry.append(func)
    return func

@register
def f1() -> int:
    return 1

@register
def f2() -> int:
    return 2

def f3() -> int:
    return 3

assert f1() == 1
assert len(registry) == 2
assert all(callable(f) for f in registry)

For a decorator to take arguments, you must add one more layer of functions.

from typing import Callable

registry = set()

def register(active=True) -> Callable[[Callable], Callable]:
    def decorate(func: Callable) -> Callable:
        if active:
            registry.add(func)
        else:
            registry.discard(func)
        return func

    return decorate

@register(active=False)
def f1() -> int:
    return 1

@register()
def f2() -> int:
    return 2

def f3() -> int:
    return 3

assert f1 not in registry
assert f2 in registry
assert f3 not in registry
# so f1 was not added because it was active=False

More complex decorators should be classes with a __call__ method.

Source

  • Python Tricks The Book, Dan Bader
  • Fluent Python, Luciano Ramalho