Decorator trong Python: hiểu và dùng đúng cách

Decorator là cách thêm chức năng cho một hàm mà không sửa mã nguồn của nó. Bạn đã gặp decorator rồi — @app.route trong Flask, @property trong lớp, @pytest.fixture khi viết kiểm thử.

Hàm là đối tượng hạng nhất

Để hiểu decorator, trước hết cần thấy rằng trong Python, hàm cũng là một giá trị như bao giá trị khác:

def chao(ten):
    return f"Xin chào, {ten}"

f = chao                    # gán hàm cho biến
print(f("Lan"))

def goi_hai_lan(ham, doi_so):
    return ham(doi_so), ham(doi_so)     # truyền hàm làm tham số

Decorator đầu tiên

import time
import functools

def do_thoi_gian(ham):
    @functools.wraps(ham)
    def ben_trong(*args, **kwargs):
        bat_dau = time.perf_counter()
        ket_qua = ham(*args, **kwargs)
        thoi_gian = time.perf_counter() - bat_dau
        print(f"{ham.__name__} chạy mất {thoi_gian:.4f} giây")
        return ket_qua
    return ben_trong

@do_thoi_gian
def tinh_toan_nang(n):
    return sum(i ** 2 for i in range(n))

tinh_toan_nang(1_000_000)

Cú pháp @do_thoi_gian tương đương với tinh_toan_nang = do_thoi_gian(tinh_toan_nang).

Đừng quên functools.wraps

Không có @functools.wraps, hàm sau khi bọc sẽ mất tên và docstring gốc:

print(tinh_toan_nang.__name__)     # 'ben_trong' nếu thiếu wraps
                                    # 'tinh_toan_nang' nếu có

Điều này gây rắc rối thật sự khi gỡ lỗi hoặc khi công cụ khác dựa vào tên hàm — chẳng hạn framework web định tuyến theo tên.

Decorator có tham số

Muốn truyền tham số cho decorator, bạn cần thêm một tầng hàm:

def thu_lai(so_lan=3, cho=1.0):
    def decorator(ham):
        @functools.wraps(ham)
        def ben_trong(*args, **kwargs):
            loi_cuoi = None
            for lan in range(so_lan):
                try:
                    return ham(*args, **kwargs)
                except Exception as e:
                    loi_cuoi = e
                    if lan < so_lan - 1:
                        print(f"Lần {lan + 1} thất bại, thử lại sau {cho}s")
                        time.sleep(cho)
            raise loi_cuoi
        return ben_trong
    return decorator

@thu_lai(so_lan=5, cho=2.0)
def goi_api():
    return requests.get("https://api.example.com").json()

Vài decorator có sẵn hữu ích

functools.cache — nhớ kết quả đã tính

from functools import cache

@cache
def fibonacci(n):
    if n < 2:
        return n
    return fibonacci(n - 1) + fibonacci(n - 2)

print(fibonacci(100))     # tức thì

Không có @cache, hàm này tính lại cùng một giá trị vô số lần và fibonacci(40) đã mất vài giây. Với cache, mỗi giá trị chỉ tính một lần.

Lưu ý: cache giữ mọi kết quả trong bộ nhớ mãi mãi. Nếu tham số đầu vào đa dạng, hãy giới hạn kích thước:

from functools import lru_cache

@lru_cache(maxsize=1024)
def tra_cuu(ma):
    ...

Chỉ dùng cache cho hàm thuần túy — cùng đầu vào luôn cho cùng đầu ra, không có tác dụng phụ. Cache một hàm đọc từ cơ sở dữ liệu sẽ khiến bạn nhận dữ liệu cũ mà không hiểu vì sao.

property — biến phương thức thành thuộc tính

class HinhChuNhat:
    def __init__(self, rong, cao):
        self._rong = rong
        self._cao = cao

    @property
    def dien_tich(self):
        return self._rong * self._cao

    @property
    def rong(self):
        return self._rong

    @rong.setter
    def rong(self, gia_tri):
        if gia_tri <= 0:
            raise ValueError("Chiều rộng phải lớn hơn 0")
        self._rong = gia_tri

hcn = HinhChuNhat(5, 3)
print(hcn.dien_tich)      # 15 — gọi như thuộc tính, không cần ngoặc
hcn.rong = -1             # ValueError

Ví dụ thực tế: kiểm tra quyền

def yeu_cau_quyen(*quyen_can):
    def decorator(ham):
        @functools.wraps(ham)
        def ben_trong(nguoi_dung, *args, **kwargs):
            thieu = set(quyen_can) - set(nguoi_dung.quyen)
            if thieu:
                raise PermissionError(f"Thiếu quyền: {', '.join(thieu)}")
            return ham(nguoi_dung, *args, **kwargs)
        return ben_trong
    return decorator

@yeu_cau_quyen("xoa_bai_viet")
def xoa_bai(nguoi_dung, bai_id):
    print(f"Đã xóa bài {bai_id}")

Logic kiểm tra quyền nằm gọn một chỗ, và mỗi hàm nghiệp vụ chỉ cần khai báo yêu cầu của nó. Khi quy tắc phân quyền thay đổi, bạn sửa một nơi duy nhất.

Ghi log tự động

import logging

def ghi_log(ham):
    @functools.wraps(ham)
    def ben_trong(*args, **kwargs):
        logging.info("Gọi %s với %s %s", ham.__name__, args, kwargs)
        try:
            ket_qua = ham(*args, **kwargs)
            logging.info("%s trả về thành công", ham.__name__)
            return ket_qua
        except Exception:
            logging.exception("%s ném ngoại lệ", ham.__name__)
            raise
    return ben_trong

Chú ý decorator này ném lại ngoại lệ sau khi ghi log. Một decorator ghi log mà nuốt luôn lỗi sẽ che giấu vấn đề thay vì giúp bạn phát hiện.

Xếp chồng nhiều decorator

@ghi_log
@do_thoi_gian
@thu_lai(so_lan=3)
def xu_ly():
    ...

Thứ tự áp dụng từ dưới lên: thu_lai bọc hàm gốc trước, rồi do_thoi_gian bọc kết quả đó, cuối cùng ghi_log bọc ngoài cùng. Khi gọi, luồng thực thi đi từ trên xuống.

Thứ tự này có ảnh hưởng thật. Ở ví dụ trên, thời gian đo được bao gồm cả các lần thử lại. Nếu bạn đảo hai decorator, mỗi lần thử sẽ được đo riêng.

Khi nào không nên dùng decorator

Decorator giấu bớt hành vi khỏi tầm mắt người đọc. Nếu logic bạn định thêm là phần cốt lõi của hàm, hãy viết thẳng vào hàm. Decorator hợp nhất với những mối quan tâm cắt ngang: ghi log, đo hiệu năng, kiểm tra quyền, cache, thử lại — những thứ lặp lại ở nhiều nơi và không thuộc về nghiệp vụ chính.