fastapi框架系列-给fastapi的路由@上一个装饰器-函数装饰器和类装饰器实现

前言

近期有收到朋友咨询关于如何给Fastapi框架加上相关类似的权限校验的装饰器处理器。然后我说推荐你直接使用注入依赖来实现比较好，但是他说可能装饰器看起来更好看，更独立一点。这个我竟无法反驳。由于之前做FLASK的时候确实都是基于装饰器的方式来校验权限，所以该来的还是会来的，下面我们就实现一下。

函数装饰器 -实现过程的记录

步骤1 定义有参装饰器

首先我们的知道权限校验装饰器，肯定是需要携带参数的，所以我们的需要定义一个有参装饰器。

def wrapper_auth(roles: Union[str, Sequence[str]] = None,
                 permissions: Union[str, Sequence[str]] = None):
    def wrapper(func):
        @functools.wraps(func)
        def _wrapper(*args: Any, **kwargs: Any) -> Response:
            print("内部函数！", roles)
            return func(*args, **kwargs)
        return _wrapper

    return wrapper

步骤2 常规装饰到路由上

@app.get('/')
@wrapper_auth
def auttlogin():
    return ''

步骤3 验证测试-咯嘣脆的异常出现了

不出意外上面这种定义的装饰器的方式有问题，首先是或报

提示我们路由上缺少一个查询的参数名字就是func,也就是我们的装饰器上面的func，这就蛋疼了！那么再来一个处理，那我们允许func为空不就OK了吗！！！是吧！！！这样就不说话必要参数了！！！

步骤4 设置func为非必选参数

def wrapper_auth(roles: Union[str, Sequence[str]] = None,
                 permissions: Union[str, Sequence[str]] = None):
    def wrapper(func:Optional[Callable]=None):
        @functools.wraps(func)
        def _wrapper(*args: Any, **kwargs: Any) -> Response:
            print("内部函数！", roles)
            return func(*args, **kwargs)
        return _wrapper

    return wrapper

步骤5 再次验证：再次咯嘣脆的异常 图片分析路由内部源码：

它需要是转化为一个是一个可调用的对象，所以我们的方面装饰写的方式就不对！！！原来的方式:

@app.get('/')
@wrapper_auth
def auttlogin():
    return ''

修改为：

@app.get('/')
@wrapper_auth()
def auttlogin():
    return ''

这样在测试就OK了！但是还是有问题！！加入我装饰的是一个协程函数的呢？如下：

@app.get('/')
@wrapper_auth()
async def auttlogin():
    return ''

此时在测试:

好吧问题多多！！！那我就在继续判断处理，需要根据被装饰是什么来做不同的处理呗！！

步骤6 优化，同时支持同步或异步装饰

def wrapper_auth(roles: Union[str, Sequence[str]] = None,
                 permissions: Union[str, Sequence[str]] = None):


    def wrapper(func:Optional[Callable]=None):
        if asyncio.iscoroutinefunction(func):
            @functools.wraps(func)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Response:
                print("装饰协程函数")
                return await func(*args, **kwargs)
            return async_wrapper
        else:
            @functools.wraps(func)
            def _wrapper(*args: Any, **kwargs: Any) -> Response:
                print("装饰同步函数！", roles)
                return func(*args, **kwargs)
            return _wrapper

    return wrapper

然后在测试：ok!!!!再进化！！我想当做装饰器也可以当做一个依赖项！你妹d！奇葩需求真多！！！既然如果是作为依赖项的话，那说明我我们的func肯定是没值的，那么我们的可以在装饰器制作做判断处理不就可以了！！！是不是！！！！！这不简单！！！

代码示例如下：

def wrapper_auth(roles: Union[str, Sequence[str]] = None,
                 permissions: Union[str, Sequence[str]] = None):


    def wrapper(func:Optional[Callable]=None):
        if func is None:
            print("被当做依赖项处理！！！")
            return True
        if asyncio.iscoroutinefunction(func):
            @functools.wraps(func)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Response:
                print("装饰协程函数")
                return await func(*args, **kwargs)
            return async_wrapper
        else:
            @functools.wraps(func)
            def _wrapper(*args: Any, **kwargs: Any) -> Response:
                print("装饰同步函数！", roles)
                return func(*args, **kwargs)
            return _wrapper

    return wrapper


# @app.get('/')
# @wrapper_auth()
# async def auttlogin():
#     return ''


@app.get('/')
@wrapper_auth()
async def auttlogin(asuv=Depends(wrapper_auth)):
    return ''

但是还是此时会把func当做依赖参数进行校验了！如下错误所示：

{"detail":[{"loc":["query","args"],"msg":"field required","type":"value_error.missing"},{"loc":["query","kwargs"],"msg":"field required","type":"value_error.missing"}]}

！！好吧！！！和我们的上面一样，值需要传入是一个对象就可以了！

步骤7 再改最终函数装饰器代码示例如下：

def wrapper_auth(roles: Union[str, Sequence[str]] = None,
                 permissions: Union[str, Sequence[str]] = None):
    def wrapper(func: Optional[Callable] = None):
        async def back_request(request: Request) -> Request:
            print("依赖项内部执行逻辑处理！！！可以在这里做权限校验逻辑")
            return request
        if func is None:
            print("被当做依赖项处理！！！需要返回是一个callable object！不然会raise TypeError('{!r} is not a callable object'.format(obj))")
            return back_request
        if asyncio.iscoroutinefunction(func):
            @functools.wraps(func)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Response:
                print("装饰协程函数！！可以在这里做权限校验逻辑", roles, args, kwargs)
                return await func(*args, **kwargs)

            return async_wrapper
        else:
            @functools.wraps(func)
            def _wrapper(*args: Any, **kwargs: Any) -> Response:
                print("装饰同步函数！！可以在这里做权限校验逻辑", roles)
                return func(*args, **kwargs)

            return _wrapper

使用示例如图示：

类方式装饰器

使用类方式定义的话，其实可以发挥的空间更大，可以自定义内部很多关于权限校验一些细节方法！

参考示例代码如下：

class AuthWrapper():

    def __init__(self,
                 roles: Union[str, Sequence[str]] = None,
                 permissions: Union[str, Sequence[str]] = None,
                 *args, **kwargs):
        self.roles = roles
        self.permissions = permissions
        self.args = args
        self.kwargs = kwargs

    def get_user_roles(self):
        pass

    def get_user_permissions(self):
        pass

    def __call__(self, func: Optional[Callable] = None):

        # def back_request(request: Request)->Request:
        #     print("依赖项内部执行逻辑处理！！！可以在这里做权限校验逻辑")
        #     return request
        async def back_request(request: Request) -> Request:
            print("依赖项内部执行逻辑处理！！！可以在这里做权限校验逻辑")
            return request

        if func is None:
            print("被当做依赖项处理！！！需要返回是一个callable object！不然会raise TypeError('{!r} is not a callable object'.format(obj))")
            return back_request
        if asyncio.iscoroutinefunction(func):
            @functools.wraps(func)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Response:
                print("装饰协程函数！！可以在这里做权限校验逻辑", self.roles, args, kwargs)

                if self.roles and self.get_user_roles not  in self.roles:
                    raise Exception("您当前的角色不允许访问滴！！！")
                if self.permissions and self.get_user_permissions not  in self.roles:
                    raise Exception("您当前的没有这个权限标记滴！！！！！！")
                try:
                    # 查当前被装饰的函数有哪些参数值信息  # 查验当前函数有哪些参数值得传入
                    sig = inspect.signature(func)
                    for idx, parameter in enumerate(sig.parameters.values()):
                        print(idx, parameter, parameter.name)
                    return await func(*args, **kwargs)
                finally:
                    pass

            return async_wrapper
        else:
            @functools.wraps(func)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Response:
                if self.roles and self.get_user_roles not in self.roles:
                    raise Exception("您当前的角色不允许访问滴！！！")
                if self.permissions and self.get_user_permissions not in self.roles:
                    raise Exception("您当前的没有这个权限标记滴！！！！！！")
                print("装饰同步函数！！可以在这里做权限校验逻辑", self.roles)
                return func(*args, **kwargs)

            return sync_wrapper


@app.get('/')
async def auttlogin(asda=Depends(AuthWrapper)):
    print(asda)
    return ''


def text_asd():
    return True


@app.get('/w')
@AuthWrapper()
async def auttlogin(asda=Depends(text_asd)):
    return ''


@app.get('/w2')
@AuthWrapper(roles=['admin', 'vip', 'novip'], permissions=['w', 'r'])
def auttlogin(asda=Depends(text_asd)):
    return ''

@app.get('/w3')
@AuthWrapper(roles=['vip'], permissions=['w', 'r'])
def auttlogin(asda=Depends(text_asd)):
    return ''

@app.get('/w4')
@AuthWrapper(roles=['superadmin'], permissions=['w', 'r'])
def auttlogin(asda=Depends(text_asd)):
    return ''

代码解析说明：

• AuthWrapper中可以传入具体的关于权限校验的规则，比如这个接口需要角色是什么样的角色：
  • 比如@app.get('/w2')就需要内置roles=['admin', 'vip', 'novip']这三个角色才可以访问， 如果不是则直接抛出异常，而且还需要permissions=['w', 'r']
• 如果类被当做依赖项处理，和函数处理逻辑类似，需要避免两个问题
  • 1：避免类中call的参数被当做查询参数进行依赖校验检测
  • 2：避免solved = await run_in_threadpool(call, **sub_values)中的TypeError: run_in_threadpool() got multiple values for argument 'func'错误。主要是需要把我们的依赖注入中像转为为一个对象。

被当做依赖项处理的时候：错误的示范为：

错误1：

@app.get('/')
async def auttlogin(asda=Depends(AuthWrapper):
    print(asda)
    return ''

上述-会引发call的参数被当做查询参数进行依赖校验检测

错误2：

@app.get('/')
async def auttlogin(asda=Depends(AuthWrapper()):
    print(asda)
    return ''

上述-TypeError: run_in_threadpool() got multiple values for argument 'func'错误

正确方式：

@app.get('/')
async def auttlogin(asda=Depends(AuthWrapper(roles=['admin', 'vip', 'novip'], permissions=['w', 'r']))):
    print(asda)
    return ''
或

@app.get('/')
async def auttlogin(asda=Depends(AuthWrapper()())):
    print(asda)
    return ''

从上面我们也可以得到一个关于依赖注入的时候如何自定义传参的问题的解决方案！！！ 如下示例：

@app.get('/')
async def auttlogin(asda=Depends(AuthWrapper(roles=['admin', 'vip', 'novip'], permissions=['w', 'r'])())):
    print(asda)
    return ''

以上仅仅是个人结合自己的实际需求，做学习的实践笔记！如有笔误！欢迎批评指正！感谢各位大佬！

结尾

END

简书：www.jianshu.com/u/d6960089b…

掘金：juejin.cn/user/296393…

公众号：微信搜【程序员小钟同学】

小钟同学 | 文 【欢迎一起学习交流】| QQ：308711822