How to add a custom decorator to a FastAPI route?

问题

I want to add an auth_required decorator to my endpoints. (Please consider that this question is about decorators, not middleware)

So a simple decorator looks like this:

def auth_required(func):
    def wrapper(*args, **kwargs):
        if user_ctx.get() is None:
            raise HTTPException(...)
        return func(*args, **kwargs)
    return wrapper

So there are 2 usages:

@auth_required
@router.post(...)

@router.post(...)
@auth_required

The first way doesn't work because router.post creates a router that saved into self.routes of APIRouter object. The second way doesn't work because it fails to verify pydantic object. For any request model, it says missing args, missing kwargs.

So my question is - how can I add any decorators to FastAPI endpoints? Should I get into router.routes and modify the existing endpoint? Or use some functools.wraps like functions?

回答1:

How can I add any decorators to FastAPI endpoints?

As you said, you need to use @functools.wraps(...)--(PyDoc) decorator as,

from functools import wraps

from fastapi import FastAPI
from pydantic import BaseModel


class SampleModel(BaseModel):
    name: str
    age: int


app = FastAPI()


def auth_required(func):
    @wraps(func)
    async def wrapper(*args, **kwargs):
        return await func(*args, **kwargs)

    return wrapper


@app.post("/")
@auth_required # Custom decorator
async def root(payload: SampleModel):
    return {"message": "Hello World", "payload": payload}

The main caveat of this method is that you can't access the request object in the wrapper and I assume it is your primary intention.

If you need to access the request, you must add the argument to the router function as,

from fastapi import Request


@app.post("/")
@auth_required  # Custom decorator
async def root(request: Request, payload: SampleModel):
    return {"message": "Hello World", "payload": payload}

I am not sure what's wrong with the FastAPI middleware, after all, the @app.middleware(...) is also a decorator.

回答2:

Here is how you can use a decorator that adds extra parameters to the route handler:

from fastapi import FastAPI, Request
from pydantic import BaseModel


class SampleModel(BaseModel):
    name: str
    age: int


app = FastAPI()

def do_something_with_request_object(request: Request):
    print(request)

def auth_required(handler):
    async def wrapper(request: Request, *args, **kwargs):
        do_something_with_request_object(request)
        return await handler(*args, **kwargs)

    # Fix signature of wrapper
    import inspect
    wrapper.__signature__ = inspect.Signature(
        parameters = [
            # Use all parameters from handler
            *inspect.signature(handler).parameters.values(),

            # Skip *args and **kwargs from wrapper parameters:
            *filter(
                lambda p: p.kind not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD),
                inspect.signature(wrapper).parameters.values()
            )
        ],
        return_annotation = inspect.signature(handler).return_annotation,
    )

    return wrapper


@app.post("/")
@auth_required # Custom decorator
async def root(payload: SampleModel):
    return {"message": f"Hello {payload.name}, {payload.age} years old!"}

来源：https://stackoverflow.com/questions/64497615/how-to-add-a-custom-decorator-to-a-fastapi-route

标签

python

python-decorators

fastapi

pydantic