如何编写Python装饰器实现API认证令牌自动刷新？

• 用过装饰器但从未自行编写过，查阅资料后仍有困惑
• 调用第三方API需携带认证令牌，令牌过期时接口返回HTTP 401状态码，此时需重新获取令牌
• 当前所有API调用函数中都通过if语句判断401并刷新令牌，认为该场景适合用装饰器优化，但不知如何实现
• 询问该问题是否适合在此平台提问

现有代码：

def get_token(self) -> str:
    r = self.session.post(
        BASE_URL + "authentication/signin",
        data=dumps({"name": self.username, "password": self.password}),
        verify=False,
    )

    if not r.ok:
        raise Exception()
    elif "token" not in (response := r.json()):
        raise Exception("Token not found in response")

    return response["token"]

def get_device(self, ip: str) -> dict[str, Any]:
    data = dumps({"ipAddress": ip.strip()})

    if (
        r := self.session.post(BASE_URL + "devices/filter", data=data)
    ).status_code == 401:
        self.session.headers.update({"X-AUTH-TOKEN": self.get_token()})
        r = self.session.post(BASE_URL + "devices/filter", data=data)

    if not r.ok:
        raise Exception()
    elif "content" not in (response := r.json()):
        raise Exception()

    return response["content"][0]

问题适配性 #

这个问题完全适合在此平台提问，Python装饰器的场景化应用、API请求的代码优化都是平台常见的技术问题范畴。

装饰器实现方案 #

你可以编写一个装饰器，把「检查401、刷新令牌、重新请求」的重复逻辑抽离出来，让API调用函数只专注于核心业务：

1. 编写基础版装饰器 #

因为你的函数都是类的实例方法，装饰器需要兼容实例的上下文，实现如下：

import functools

def handle_token_expired(func):
    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        # 首次执行API请求
        resp = func(self, *args, **kwargs)
        # 判断是否令牌过期
        if resp.status_code == 401:
            # 刷新令牌并更新请求头
            new_token = self.get_token()
            self.session.headers.update({"X-AUTH-TOKEN": new_token})
            # 重新发起请求
            resp = func(self, *args, **kwargs)
        # 校验最终请求状态
        if not resp.ok:
            raise Exception(f"API请求失败，状态码: {resp.status_code}")
        return resp
    return wrapper

2. 改造原有API调用函数 #

移除get_device里的401判断逻辑，只保留核心请求代码，再给函数加上装饰器：

def get_token(self) -> str:
    r = self.session.post(
        BASE_URL + "authentication/signin",
        data=dumps({"name": self.username, "password": self.password}),
        verify=False,
    )

    if not r.ok:
        raise Exception(f"获取令牌失败，状态码: {r.status_code}")
    elif "token" not in (response := r.json()):
        raise Exception("Token not found in response")

    return response["token"]

@handle_token_expired
def get_device(self, ip: str):
    data = dumps({"ipAddress": ip.strip()})
    # 直接返回请求对象，交给装饰器处理后续逻辑
    return self.session.post(BASE_URL + "devices/filter", data=data)

3. 进阶：带响应解析的装饰器 #

如果需要统一处理JSON解析和字段校验，可以扩展装饰器，让它直接返回业务数据：

def handle_token_expired_with_parse(func):
    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        resp = func(self, *args, **kwargs)
        if resp.status_code == 401:
            new_token = self.get_token()
            self.session.headers.update({"X-AUTH-TOKEN": new_token})
            resp = func(self, *args, **kwargs)
        
        if not resp.ok:
            raise Exception(f"API请求失败，状态码: {resp.status_code}")
        
        resp_json = resp.json()
        if "content" not in resp_json:
            raise Exception("响应中缺少content字段")
        return resp_json["content"][0]
    return wrapper

# 此时get_device可以简化到极致
@handle_token_expired_with_parse
def get_device(self, ip: str):
    data = dumps({"ipAddress": ip.strip()})
    return self.session.post(BASE_URL + "devices/filter", data=data)

4. 装饰器的优势 #

• 代码复用：所有需要处理令牌过期的API函数，只需加个装饰器，不用重复写401判断逻辑
• 逻辑分离：业务请求和令牌刷新逻辑拆分，代码更清晰易维护
• 扩展性强：后续要处理其他异常状态码，只需修改装饰器即可

内容的提问来源于stack exchange，提问作者Net-worker