【Python 装饰器】实战:从计时器到登录验证
从最简单的装饰器开始写三个例子逐步升级最后写出 Flask JWT 接口里的login_required。每个例子都能直接运行建议边看边敲。引言回顾通用模板 所有装饰器都是这个结构的变体先记住它defmy_decorator(func):defwrapper(*args,**kwargs):# 前置逻辑原函数执行前resultfunc(*args,**kwargs)# 后置逻辑原函数执行后returnresultreturnwrapper 三个位置可以填代码执行前、执行后、以及决定要不要执行原函数。下面三个例子会分别用到这三个位置。一、例子一计时器测量函数执行时间 最实用的入门装饰器。你想知道一个函数跑了多久但不想在每个函数里都加计时代码。1.1 先看不用装饰器的写法importtimedefslow_function():time.sleep(1)print(执行完毕)starttime.time()slow_function()endtime.time()print(f耗时{end-start:.2f}秒) 如果有 10 个函数都要计时每个都写一遍start time.time()…end time.time()需要重复写10次。1.2 用装饰器的写法importtimedeftimer(func):defwrapper(*args,**kwargs):starttime.time()# 前置记录开始时间resultfunc(*args,**kwargs)# 执行原函数endtime.time()# 后置记录结束时间print(f{func.__name__}耗时{end-start:.2f}秒)returnresultreturnwrapper 对照通用模板前置逻辑是记录开始时间后置逻辑是记录结束时间并打印差值。func.__name__拿到原函数的名字这样打印出来就知道是哪个函数。使用timerdefslow_function():time.sleep(1)print(执行完毕)timerdeffast_function():print(秒完)slow_function()fast_function()执行完毕 slow_function 耗时1.00 秒 秒完 fast_function 耗时0.00 秒timer往上一挂就行想给哪个函数计时就挂哪个不需要改函数本身的代码。不想计时了把timer删掉就行函数本身一行都不用动。二、例子二日志记录记录谁调用了什么 升级一下写一个记录函数调用信息的装饰器函数名、参数、返回值。deflog(func):defwrapper(*args,**kwargs):print(f调用{func.__name__}参数args{args}, kwargs{kwargs})resultfunc(*args,**kwargs)print(f{func.__name__}返回{result})returnresultreturnwrapper 使用logdefadd(a,b):returnablogdefgreet(name,greeting你好):returnf{greeting}{name}add(1,2)greet(zhangsan)greet(lisi,greeting嗨)调用 add参数args(1,2),kwargs{}add返回3 调用 greet参数args(zhangsan,),kwargs{}greet 返回你好zhangsan 调用 greet参数args(lisi,),kwargs{greeting:嗨}greet 返回嗨lisi 这里能看到*args和**kwargs的实际效果add(1, 2)的两个参数被args捕获为元组(1, 2)greet(lisi, greeting嗨)的关键字参数被kwargs捕获为字典{greeting: 嗨}。不管原函数长什么样装饰器都能通用。三、例子三权限检查决定让不让执行 前两个例子都是前后加点逻辑但原函数一定会执行。这个例子不一样装饰器要决定原函数能不能执行。这正是login_required的核心模式。 先写一个简单版检查用户是不是管理员是才让执行不是就拒绝。3.1 权限检查装饰器实现current_user{name:zhangsan,role:user}# 模拟当前用户defadmin_required(func):defwrapper(*args,**kwargs):ifcurrent_user[role]!admin:print(f权限不足{current_user[name]}不是管理员)returnNone# 不执行原函数直接返回returnfunc(*args,**kwargs)# 是管理员正常执行returnwrapper 和前两个例子的关键区别wrapper里有一个if判断不满足条件就return了根本不会调用func()。原函数被拦在了外面。3.2 删除用户admin_requireddefdelete_user(user_id):print(f已删除用户{user_id})returnTrue如果是普通用户即zhangsandelete_user(123)# 权限不足zhangsan 不是管理员换成管理员current_user{name:admin,role:admin}delete_user(123)已删除用户123 是不是和login_required的模式很像了login_required做的就是同样的事检查 Token 是否合法合法才让请求进入接口函数不合法就直接返回 401。四、wraps修复装饰器的一个副作用 在进入login_required之前还需要解决一个问题。4.1 默认情况 看看装饰后函数的名字timerdefslow_function():这是一个慢函数time.sleep(1)print(slow_function.__name__)print(slow_function.__doc__)wrapper None 名字变成了wrapper文档字符串也丢了。因为slow_function现在指向的是wrapper函数而不是原来的slow_function。4.2 为什么这不只是好看的问题 这在 Flask 里会引发真实的 bugFlask 用__name__作为路由的 endpoint 名如果两个视图函数装饰后名字都变成wrapperFlask 会抛出AssertionError提示 endpoint 重复注册。4.3 用 wraps 修复 从functools导入wraps加在wrapper上面fromfunctoolsimportwrapsdeftimer(func):wraps(func)# 加这一行defwrapper(*args,**kwargs):starttime.time()resultfunc(*args,**kwargs)endtime.time()print(f{func.__name__}耗时{end-start:.2f}秒)returnresultreturnwrappertimerdefslow_function():这是一个慢函数time.sleep(1)print(slow_function.__name__)# slow_function ← 名字保住了print(slow_function.__doc__)# 这是一个慢函数 ← 文档也保住了wraps(func)把原函数的名字、文档字符串等属性复制到wrapper上。写装饰器时永远加上wraps(func)这是最佳实践。4.4 更新后的通用模板fromfunctoolsimportwrapsdefmy_decorator(func):wraps(func)defwrapper(*args,**kwargs):# 前置逻辑resultfunc(*args,**kwargs)# 后置逻辑returnresultreturnwrapper五、终极实战写出 login_required 现在具备了所有知识可以理解Flask JWT 登录接口里的login_required了。先回顾一下它做的事从请求头取Authorization: Bearer token取出 Token用密钥验签验签通过 → 把用户信息挂到request.user放行验签失败 → 直接返回 401不执行接口函数 这就是例子三权限检查的模式判断条件决定让不让执行。只不过判断条件从是不是管理员变成了Token 是否合法。fromfunctoolsimportwrapsfromflaskimportrequest,jsonifyimportjwtdeflogin_required(func):wraps(func)defwrapper(*args,**kwargs):# 第一步从请求头取 Tokenauth_headerrequest.headers.get(Authorization,)ifnotauth_header.startswith(Bearer ):returnjsonify({error:缺少 Token}),401tokenauth_header.split( )[1]# 第二步验签try:datajwt.decode(token,SECRET_KEY,algorithms[HS256])exceptjwt.ExpiredSignatureError:returnjsonify({error:Token 已过期}),401exceptjwt.InvalidTokenError:returnjsonify({error:无效的 Token}),401# 第三步验签通过放行request.userdatareturnfunc(*args,**kwargs)returnwrapper 对照通用模板看模板位置login_required 里做了什么前置逻辑取 Token、验签决定是否执行原函数验签失败 → return 401不调用func()func(*args, **kwargs)验签通过 → 执行接口函数如profile()后置逻辑无接口返回什么就返回什么使用时app.route(/profile)login_requireddefprofile():returnjsonify({name:request.user[username]})login_required等价于profile login_required(profile)。请求进来时先执行wrapper验 Token通过了才执行原来的profile()。 和你在例子三里写的admin_required本质上完全一样只是判断条件更复杂验签而不是查角色、返回值更规范Flask 的 JSON 响应而不是 print。六、装饰器的执行顺序 Flask 接口上经常挂多个装饰器app.route(/profile)login_requireddefprofile():... 多个叠在一起Python 从最靠近函数的那个开始一层层往上包# Python 实际执行的等价形式profilelogin_required(profile)# 第一步login_required 先包裹app.route(/profile)(profile)# 第二步route 再把结果注册到路由表 请求来的时候调用方向正好反过来外层先执行逐层往里走。app.route负责把这个包好的函数注册进路由表。当请求匹配到/profile时Flask 调用注册的函数也就是最外层的login_required的 wrapper它先验 Token通过后才调用内层的profile。 我们用一个直观的例子来理解可能更为方便左边是定义阶段Python 从最靠近函数的login_required开始包裹再往外被app.route包裹像穿衣服贴身的先穿。右边是调用阶段与请求进来时方向相反先经过最外层的路由匹配再经过login_required验 Token最后才到达真正的profile() 函数就像脱衣服先脱外套。七、总结 回顾一下这篇做了什么从同一个模板出发填不同的逻辑写出了三种装饰器例子在模板里填了什么学到的新东西计时器func()前后记录时间装饰器的基本写法日志记录func()前后打印参数和返回值*args, **kwargs让装饰器适配任意函数权限检查func()前面加if不满足条件就不调用装饰器可以拦截请求不是非得调用原函数