Dekoratory w Pythonie: @funkcje, argumenty i functools.wraps

Dekorator to funkcja, która opakowuje funkcję

To zdanie brzmi abstrakcyjnie, ale mechanika jest prosta. Dekorator przyjmuje funkcję i zwraca funkcję. Zwracana funkcja zwykle wywołuje oryginał, a wokół tego wywołania dokłada jakieś dodatkowe zachowanie.

Najkrótszy możliwy przykład:

def shout(func):
    def wrapper():
        result = func()
        return result.upper()
    return wrapper

def greet():
    return "hello"

greet = shout(greet)
print(greet())

shout to dekorator. Przyjmuje funkcję (greet), buduje nową funkcję (wrapper), która wywołuje oryginał i zamienia wynik na wielkie litery, a potem ją zwraca. Przypisanie greet = shout(greet) podmienia oryginał na wersję opakowaną.

Ten wzorzec ponownego przypisania jest tak częsty, że Python dał mu osobną składnię.

@ to lukier składniowy dla ponownego przypisania

@name w linii nad def jest równoważne name = name(...) tuż po zdefiniowaniu funkcji:

def shout(func):
    def wrapper():
        return func().upper()
    return wrapper

@shout
def greet():
    return "hello"

print(greet())

@shout czyta się jako "zastosuj dekorator shout do tej funkcji". Python wykonuje greet = shout(greet) zaraz po def: ta sama mechanika co wcześniej, mniej pisania.

Gdy widzisz @name, w myślach zamień to na function = name(function). Tylko tyle znaczy ta składnia.

Obsługa argumentów

Większość funkcji przyjmuje argumenty. Użyteczny dekorator przekazuje je dalej. Idiom to *args, **kwargs, czyli pythonowy sposób na przyjęcie dowolnych argumentów, bo funkcję opakowującą nie powinno obchodzić, czego oczekuje opakowana funkcja:

def shout(func):
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)
        return result.upper()
    return wrapper

@shout
def greet(name, punctuation="!"):
    return f"hello, {name}{punctuation}"

print(greet("rosa"))
print(greet("boris", punctuation="."))

*args przechwytuje wszystkie argumenty pozycyjne. **kwargs przechwytuje wszystkie argumenty nazwane. Funkcja opakowująca przekazuje wszystko bez zmian do opakowanej funkcji, a potem wykonuje dodatkową pracę, dla której istnieje dekorator: tutaj zamienia wynik na wielkie litery.

Tak wygląda większość prawdziwych dekoratorów.

Bardziej przydatny przykład: pomiar czasu

Wypisz, ile trwa wykonanie funkcji:

import time

def timed(func):
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        result = func(*args, **kwargs)
        elapsed = time.perf_counter() - start
        print(f"{func.__name__} took {elapsed*1000:.2f} ms")
        return result
    return wrapper

@timed
def slow_sum(n):
    total = 0
    for i in range(n):
        total += i
    return total

print(slow_sum(1_000_000))
print(slow_sum(5_000_000))

Ten wzorzec (zrób coś przed wywołaniem i coś po nim) to to, co ostatecznie robi większość dekoratorów. Logowanie, sprawdzanie uprawnień, ponawianie i walidacja wejścia mają ten sam kształt.

Zachowanie oryginalnej tożsamości: functools.wraps

Udekorowanie funkcji ją podmienia, więc opakowana funkcja traci swoje oryginalne atrybuty __name__ i __doc__:

def shout(func):
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs).upper()
    return wrapper

@shout
def greet(name):
    """Return a personalised greeting."""
    return f"hello, {name}"

print(greet.__name__)
print(greet.__doc__)

greet.__name__ to teraz "wrapper", a docstring zniknął. To psuje help(), tracebacki i każde narzędzie, które analizuje funkcję.

Poprawka zajmuje jedną linię: @functools.wraps(func) na funkcji wewnętrznej kopiuje metadane.

from functools import wraps

def shout(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs).upper()
    return wrapper

@shout
def greet(name):
    """Return a personalised greeting."""
    return f"hello, {name}"

print(greet.__name__)
print(greet.__doc__)

Zawsze dodawaj @wraps(func) do funkcji wewnętrznej. Nic nie kosztuje, a oszczędza zaskakujących sesji debugowania w przyszłości.

Dekoratory z argumentami

Czasem sam dekorator potrzebuje konfiguracji: "ponów tę funkcję maksymalnie 3 razy", "loguj na poziomie DEBUG". To oznacza jeszcze jedną warstwę zagnieżdżenia: zewnętrzną funkcję, która przyjmuje argumenty i zwraca dekorator.

from functools import wraps

def repeat(times):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            for _ in range(times):
                result = func(*args, **kwargs)
            return result
        return wrapper
    return decorator

@repeat(times=3)
def say_hi(name):
    print(f"Hi, {name}")

say_hi("Rosa")

Trzy warstwy brzmią groźnie. Czytaj od zewnątrz:

1. repeat(times=3) to wywołanie funkcji. Zwraca decorator.
2. decorator to właściwy dekorator: przyjmuje funkcję i zwraca opakowaną.
3. wrapper to opakowana funkcja, która działa w chwili wywołania.

Na tym kształcie opierają się @retry(times=5), @cache(maxsize=100) i dekoratory frameworków, takie jak @app.route("/users"). Gdy zobaczysz trójwarstwowy wzorzec, całą tę rodzinę czyta się tak samo.

Nakładanie dekoratorów

Na jedną funkcję możesz nałożyć więcej niż jeden dekorator. Nakładają się od dołu do góry: ten najbliżej def działa pierwszy:

from functools import wraps

def add_exclaim(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs) + "!"
    return wrapper

def shout(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs).upper()
    return wrapper

@shout
@add_exclaim
def greet(name):
    return f"hi, {name}"

print(greet("rosa"))

add_exclaim opakowuje pierwszy i dodaje !. Potem shout opakowuje to wszystko i zamienia na wielkie litery. Wynik to HI, ROSA!.

Kolejność ma znaczenie. Po odwróceniu kolejności dostaniesz HI, ROSA! z wykrzyknikiem dodanym po zamianie na wielkie litery. Tutaj wygląda to identycznie, ale wyobraź sobie dekorator formatujący JSON: uruchomienie go przed albo po dekoratorze, który loguje wejście, może dać bardzo różne wyniki.

Wbudowane dekoratory, które zobaczysz

Python i jego biblioteka standardowa mają garść dekoratorów, które spotkasz w prawdziwym kodzie:

from functools import lru_cache

class Circle:
    def __init__(self, radius):
        self.radius = radius

    @property
    def area(self):
        return 3.14159 * self.radius ** 2

    @staticmethod
    def unit_circle_area():
        return 3.14159

    @classmethod
    def from_diameter(cls, diameter):
        return cls(diameter / 2)

@lru_cache(maxsize=128)
def fib(n):
    if n < 2:
        return n
    return fib(n - 1) + fib(n - 2)

c = Circle(5)
print(c.area)
print(Circle.unit_circle_area())
print(Circle.from_diameter(10).radius)
print(fib(50))

• @property zamienia metodę w obliczany atrybut.
• @staticmethod oznacza metodę, która nie używa self ani cls.
• @classmethod dostaje klasę jako cls zamiast instancji, co świetnie sprawdza się w alternatywnych konstruktorach.
• @functools.lru_cache zapamiętuje wyniki, więc powtarzane wywołania z tymi samymi argumentami trafiają do cache'u.

Dekoratory frameworków (@app.route, @pytest.fixture, @dataclass) działają na tej samej zasadzie. Nic specjalnego: to tylko funkcje, które opakowują funkcje.

Kiedy pisać dekorator, a kiedy nie

Napisz dekorator, gdy chcesz zastosować to samo zachowanie do wielu funkcji: pomiar czasu, logowanie, ponawianie, sprawdzanie uprawnień. Cały sens polega na tym, że to zachowanie zostaje poza ciałem funkcji.

Odpuść dekorator, gdy:

• Zachowanie należy do jednej konkretnej funkcji. Umieść je w tej funkcji.
• Potrzebujesz go tylko do testów. Fixture albo parametr będzie czytelniejszy.
• Kusi cię nałożenie czterech czy pięciu dekoratorów. Wtedy przepływ sterowania chowa się w łańcuchu dekoratorów i czytelnik musi rozwinąć każdą warstwę, żeby zobaczyć, co naprawdę się wykonuje. Zwykła funkcja pomocnicza może czytać się lepiej.

Dekoratory to ostre narzędzie. Dobrze użyte utrzymują kod zgodny z DRY, a intencję oczywistą. Źle użyte ukrywają, co robi program. Decydując, czy po nie sięgnąć, stawiaj na "oczywiste".

Dalej: type hints

Dekoratory to częste miejsce, w którym spotyka się type hints w praktyce, bo funkcje wrapper często mają adnotowane sygnatury. Type hints to niewielka funkcja języka, która szybko się zwraca, i o nich jest następna strona.