如何自定义插件

Pait自带是插件并不多,但开发者可以根据自己的需求实现插件,下面以异常捕获插件为例子阐述如何制作一个插件。

下面所示代码是一个简单的API接口:

from typing import Optional
import uvicorn  # type: ignore
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.routing import Route
from pait.exceptions import TipException

from pait.app.starlette import pait
from pait import field


async def api_exception(request: Request, exc: Exception) -> JSONResponse:
    """提取异常信息, 并以响应返回"""
    if isinstance(exc, TipException):
        exc = exc.exc
    return JSONResponse({"data": str(exc)})


@pait()
async def demo(
    uid: str = field.Query.i(),
    user_name: Optional[str] = field.Query.i(default=None),
    email: Optional[str] = field.Query.i(default=None)
) -> JSONResponse:
    return JSONResponse({"uid": uid, "user_name": user_name, "email": email})


app = Starlette(routes=[Route("/api/demo", demo, methods=["GET"])])
app.add_exception_handler(Exception, api_exception)
uvicorn.run(app)
这个接口由Pait提供参数校验功能,如果调用方发起的参数有错,则会直接抛出异常并最终被starlette捕获再分发到api_exception函数处理,比如下面的请求, Pait在校验发现缺少参数uid时会抛出错误,最后被api_exception捕获并把异常返回给调用方:
➜  ~ curl http://127.0.0.1:8000/api/demo
{"data":"Can not found uid value"}
现在该接口多了一个需求,需要对该路由函数的异常的处理定制化,生成不一样的返回格式,但是api_exception是统一处理所有接口函数的异常, 它不可能为每个函数定义一个单独的函数处理,这时候可以定制一个捕获异常的插件来解决这个问题,如下是一个单独针对这个接口定制的插件:
from typing import Any, Dict
from pait.plugin.base import BasePlugin
from pait.model.core import PaitCoreModel
from pydantic import ValidationError
from pait.exceptions import PaitBaseException


class DemoExceptionPlugin(BasePlugin):
    is_pre_core: bool = True

    @classmethod
    def pre_check(cls, pait_core_model: "PaitCoreModel", kwargs: Dict) -> Dict:
        if pait_core_model.func.__name__ != "demo":
            raise RuntimeError(f"The {cls.__name__} is only used for demo func")
        return super().cls_hook_by_core_model(pait_core_model, kwargs)

    async def __call__(self, *args: Any, **kwargs: Any) -> Any:
        try:
            return await self.call_next(args, kwargs)
        except (ValidationError, PaitBaseException) as e:
            return JSONResponse({"plugin exc info": str(e)})
在这个示例插件中,需要注意的有几个地方:

  • 0.由于该路由函数是async的,所以__call__方法需要加上async。
  • 1.第9行的is_pre_core = True是设置该插件为前置插件,这样就能拦截Pait和路由函数的异常了。
  • 2.第12行的pre_check方法会进行一些初始化的检查,该检查只会在初始化的时候运行,这个检查的逻辑是如果判定该插件并不是挂在demo函数上就会抛错, 其中pait_core_modelPait为路由函数生成的一些属性。
  • 3.第17行的__call__方法是该插件的主要处理逻辑,当有请求进来时,Pait会通过__call__方法调用插件,插件可以通过call_next来调用下一个插件, 该插件通过try...except来捕获后续所有调用段异常,如果是符合条件的异常就会被捕获,并生成不一样的响应结果。

编写完毕插件后,就可以直接使用了,对上面代码进行如下的小更改:

@pait(plugin_list=[DemoExceptionPlugin.build])
async def demo(...): pass
然后重启程序并运行同样的请求,可以发现响应结果已经变为插件自己抛出的结果:
➜  ~ curl http://127.0.0.1:8000/api/demo
{"plugin exc info":"File \"/home/so1n/demo.py\", line 48, in demo.\nerror:Can not found uid value"}

回到页面顶部