pydantic_async_decorator

pydantic_async_decorator.py

from functools import wraps
from typing import Any
from typing import Awaitable
from typing import Callable
from typing import TypeVar

from pydantic import BaseModel


class BaseRequest(BaseModel):
    field: int


class BaseResponse(BaseModel):
    field: int


T = TypeVar("T_co", bound=BaseRequest)


def async_decorator(
    func: Callable[[T, Any], Awaitable[BaseResponse]],
) -> Callable[[T, Any], Awaitable[int | BaseResponse]]:
    """
    Wraps an async function, making a wrapper that returns
    its output or *something else* (in this case, for example purposes,
    an `int`).

    Both `T` and the return type (`BaseResponse`) need to
    be concrete Pydantic types, otherwise type inference does
    not work.
    """

    @wraps(func)
    async def wrapper(req: T, *args, **kwargs):
        out = await func(req, *args, **kwargs)

        # do things
        return -1

    return wrapper