使用装饰器检查用户操作权限

假设我们使用flask进行开发,需要检查当前用户是否具备某个操作权限。

from flask import abort
from flask_login import current_user
from .models import Permission

先从最简单的情况来想。f是某个操作函数,如添加评论。那么用户想要添加评论需要具备评论权限Permission.COMMENT
检查操作权限最直白的逻辑当然是使用条件判断语句:

if not current_user.can(Permission.COMMENT):
    abort(403)
f(*args, **kwargs)

为了不在每一次调用f前都写一遍if语句,我们将这一层逻辑包含到一个新的函数里,则有:

def decorated_function(*args, **kwargs):
    if not current_user.can(Permission.COMMENT):
        abort(403)
    return f(*args, **kwargs)

除非当前用户没有操作权限,decorated_function做的事情将与f做的事情一模一样。我们想要的就是一个这样的函数。
然而,目前decorated_function只能完成和f同样的事。如果对另一个函数g想要一个同样的函数,是否就要再写一个

def decorated_function_new(*args, **kwargs):
    if not current_user.can(Permission.COMMENT):
        abort(403)
    return g(*args, **kwargs)

了呢?

decorated_function_newdecorated_function唯一的区别在于返回的值是f(*args, **kwargs)还是g(*args, **kwargs)。那么我们可以再定义一个函数,将f或者g作为参数传进去,就可以解决问题。

def decorator(f):
    def decorated_function(*args, **kwargs):
        if not current_user.can(Permission.COMMENT):
            abort(403)
        return f(*args, **kwargs)
    return decorated_function

这样,对于任何一个函数f,我们都能通过调用decorator(f)的方式,得到一个与之对应的包含了在进行操作前检查评论权限的函数。注意这只是个函数,需要调用(和传参)才能得到返回值。

decorator(f)(*args, **kwargs)

这里,函数decorator就是一个装饰器,可以在不更改f定义的前提下在代码运行期间动态地增加f的功能。

@将这个装饰器放在某个函数定义之前,就可以起到替换原函数的作用。

@decorator
def add_comment():
    '''add a comment'''
    ...

这等价于:

add_comment = decorator(add_comment)

然而因为装饰器是替换原函数,我们失去了原函数的其他一些特征。比如:

print(add_comment.__name__)
print(add_comment.__doc__)

将输出decorated_functionNone,而不是add_commentadd a comment。如果用decorator装饰了多个函数,则它们的__name__和__doc__属性都会被掩盖掉,有些依赖函数签名的代码执行就会出错。为了避免这样的事情,可以使用functools库中的wrapsf的相关属性复制进来。

from functools import wraps

...

def decorator(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if not current_user.can(Permission.COMMENT):
            abort(403)
        return f(*args, **kwargs)
    return decorated_function

现在我们已经实现了一个装饰器,但只能用它检查评论权限。我们可以定义一个装饰器工厂函数,产生检查各种权限的装饰器。

def permission_required(permission):
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            if not current_user.can(permission):
                abort(403)
            return f(*args, **kwargs)
        return decorated_function
    return decorator

...

@permission_required(Permission.WRITE)
def new_post(user):
    ...

在这颗裹了三层def的洋葱中,最内层def添加了原函数之外的逻辑,第二层deff作为参数,使得装饰器可以在其他任何函数上,第三层用于传入参数控制装饰器的行为。

如果我们希望用permission_required为检查管理员权限单独生成一个装饰器:

def admin_required(f):
    return permission_required(Permission.ADMIN)(f)

写成admin_required = permission_required(Permission.ADMIN)应该也可以,不过还是用定义函数的形式比较好吧。

你可能感兴趣的