Decorator trong Python: cú pháp @, đối số và functools.wraps

Decorator là hàm bọc một hàm khác

Câu đó nghe trừu tượng, nhưng cơ chế thì đơn giản. Một decorator nhận vào một hàm và trả ra một hàm. Hàm được trả ra thường gọi hàm gốc, với một chút hành vi bổ sung bọc quanh lời gọi đó.

Ví dụ ngắn nhất có thể:

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

def greet():
    return "hello"

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

shout là decorator. Nó nhận một hàm (greet), tạo một hàm mới (wrapper) gọi hàm gốc rồi viết hoa kết quả, và trả hàm mới đó về. Gán lại greet = shout(greet) thay hàm gốc bằng phiên bản đã được bọc.

Mẫu gán lại này phổ biến tới mức Python dành riêng cho nó một cú pháp.

Dấu @ là cách viết gọn của phép gán lại

@name ở dòng ngay trên def tương đương với name = name(...) ngay sau khi hàm được định nghĩa:

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

@shout
def greet():
    return "hello"

print(greet())

@shout được đọc là "áp dụng decorator shout cho hàm này". Python chạy greet = shout(greet) ngay sau def: cơ chế y như trước, chỉ đỡ phải gõ.

Mỗi khi thấy @name, hãy thầm thay nó bằng function = name(function). Cú pháp chỉ có nghĩa như vậy.

Xử lý đối số

Phần lớn hàm đều nhận đối số. Một decorator dùng được phải chuyển chúng qua. Cách viết quen thuộc là *args, **kwargs (cách Python nhận mọi đối số), vì hàm bọc không nên quan tâm hàm được bọc cần gì:

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 bắt mọi đối số vị trí. **kwargs bắt mọi đối số từ khóa. Hàm bọc chuyển nguyên vẹn mọi thứ sang hàm được bọc, rồi làm phần việc bổ sung mà decorator đảm nhận, ở đây là viết hoa kết quả.

Đây là hình dạng của phần lớn decorator thực tế.

Một ví dụ hữu ích hơn: đo thời gian

In ra một hàm chạy mất bao lâu:

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))

Mẫu này (chạy gì đó trước lời gọi, chạy gì đó sau lời gọi) là việc mà phần lớn decorator rốt cuộc đều làm. Ghi log, kiểm tra xác thực, thử lại và kiểm tra đầu vào đều theo cùng hình dạng đó.

Giữ lại danh tính gốc: functools.wraps

Gắn decorator cho một hàm nghĩa là thay thế hàm đó, nên hàm được bọc mất đi các thuộc tính __name__ và __doc__ gốc:

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__ giờ là "wrapper" và docstring đã biến mất. Điều đó làm hỏng help(), traceback và mọi công cụ soi xét hàm.

Cách sửa chỉ một dòng: @functools.wraps(func) đặt trên hàm bên trong sẽ sao chép siêu dữ liệu sang.

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__)

Luôn thêm @wraps(func) cho hàm bên trong. Nó không tốn gì cả và giúp bạn tránh những buổi gỡ lỗi đầy bất ngờ sau này.

Decorator có đối số

Đôi khi chính decorator cần được cấu hình: "thử lại hàm này tối đa 3 lần", "ghi log ở mức DEBUG". Khi đó cần thêm một lớp lồng nữa: một hàm bên ngoài nhận đối số và trả về một decorator.

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")

Ba lớp nghe có vẻ nhiều. Hãy đọc từ ngoài vào:

1. repeat(times=3) là một lời gọi hàm. Nó trả về decorator.
2. decorator là decorator thực sự: nó nhận một hàm và trả về hàm đã được bọc.
3. wrapper là hàm đã được bọc, chạy vào lúc gọi.

Hình dạng này là nền tảng của @retry(times=5), @cache(maxsize=100) và các decorator của framework như @app.route("/users"). Khi đã nhận ra mẫu ba lớp, bạn sẽ đọc cả họ decorator này theo cùng một cách.

Xếp chồng decorator

Bạn có thể áp dụng nhiều decorator cho một hàm. Chúng xếp chồng từ dưới lên: decorator gần def nhất được áp dụng trước:

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 bọc trước, thêm dấu !. Sau đó shout bọc lớp đó, viết hoa tất cả. Đầu ra là HI, ROSA!.

Thứ tự có ý nghĩa. Đảo ngược thứ tự cho bạn HI, ROSA! với dấu chấm than được thêm sau khi viết hoa: nhìn thì giống hệt ở đây, nhưng hãy tưởng tượng một decorator định dạng JSON: chạy nó trước hay sau một decorator ghi log đầu vào có thể cho kết quả rất khác nhau.

Những decorator có sẵn bạn sẽ gặp

Python và thư viện chuẩn có sẵn một vài decorator mà bạn sẽ gặp trong code thực tế:

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 biến một phương thức thành thuộc tính được tính toán.
• @staticmethod đánh dấu một phương thức không dùng self hay cls.
• @classmethod nhận lớp vào làm cls thay vì một instance, rất hợp cho các hàm khởi tạo thay thế.
• @functools.lru_cache ghi nhớ kết quả, nên những lần gọi lặp lại với cùng đối số sẽ lấy từ cache.

Các decorator của framework (@app.route, @pytest.fixture, @dataclass) dùng đúng cơ chế này. Không có gì đặc biệt: chỉ là những hàm bọc hàm.

Khi nào nên viết, khi nào không

Hãy viết decorator khi bạn muốn áp dụng cùng một hành vi cho nhiều hàm: đo thời gian, ghi log, thử lại, kiểm tra quyền. Mục đích chính là để hành vi đó nằm ngoài thân hàm.

Đừng dùng decorator khi:

• Hành vi đó chỉ thuộc về một hàm cụ thể. Hãy đặt nó trong hàm.
• Bạn chỉ cần nó để kiểm thử. Một fixture hoặc một tham số sẽ rõ ràng hơn.
• Bạn bị cám dỗ xếp chồng bốn năm cái. Tới lúc đó, luồng điều khiển bị giấu trong chuỗi decorator: người đọc phải gỡ từng lớp mới thấy thứ gì thật sự chạy. Một hàm hỗ trợ đơn giản có thể dễ đọc hơn.

Decorator là công cụ sắc bén. Dùng tốt, chúng giữ code không lặp lại (DRY) và ý định rõ ràng. Dùng tệ, chúng che giấu những gì chương trình đang làm. Khi phân vân có nên dùng hay không, hãy nghiêng về phía "rõ ràng".

Tiếp theo: type hint

Decorator là nơi bạn hay gặp type hint trong code thực tế: các hàm wrapper thường chú thích kiểu cho chữ ký của chúng. Type hint là một tính năng nhỏ nhưng nhanh chóng mang lại lợi ích, và đó là chủ đề tiếp theo.