litestar依赖注入

15 阅读6分钟

litestar依赖注入

本文内容基于《官方文档-Dependency Injection》的部分进行制作,对其内容大幅进行缩减,使其能适合新手入门,如何想充分学习litestar可以直接查看官方文档

Litestar 拥有简单但强大的依赖注入系统,允许在应用程序的所有层级声明依赖项:

from litestar import Controller, Router, Litestar, get
from litestar.di import Provide

async def bool_fn() -> bool: ...
async def dict_fn() -> dict: ...
async def list_fn() -> list: ...
async def int_fn() -> int: ...

class MyController(Controller):
    path = "/controller"
    # 在控制器层级声明依赖
    dependencies = {"controller_dependency": Provide(list_fn)}

    # 在路由处理函数层级声明依赖
    @get(path="/handler", dependencies={"local_dependency": Provide(int_fn)})
    def my_route_handler(
        self,
        app_dependency: bool,
        router_dependency: dict,
        controller_dependency: list,
        local_dependency: int,
    ) -> None: ...

# 在路由层级声明依赖
my_router = Router(
    path="/router",
    dependencies={"router_dependency": Provide(dict_fn)},
    route_handlers=[MyController],
)

# 在应用层级声明依赖
app = Litestar(
    route_handlers=[my_router],
    dependencies={"app_dependency": Provide(bool_fn)}
)

注意:Litestar 在运行时需要注入的类型,这可能与类型检查工具(如 linter)推荐使用 TYPE_CHECKING 的规则冲突。

依赖项可以是可调用对象(同步 / 异步函数、方法或实现了 __call__ 方法的类实例)或类,这些对象需要通过 Provide 类包装。

1. 同步与异步可调用对象

Litestar 支持同步和异步可调用对象。需要注意的是,同步函数若执行阻塞操作(如 I/O 或计算密集型任务)可能阻塞事件循环主线程,进而影响整个应用。
为缓解此问题,可将 sync_to_thread 参数设为 True,使函数在线程池中运行。
若同步函数是非阻塞的,设为 False 可告知 Litestar 该函数无需线程池。
若未显式设置 sync_to_thread,Litestar 将发出警告。

2. 前置条件与作用域

依赖注入的前置条件:

  • 依赖项必须是可调用对象。
  • 依赖项可接收 kwargsself 参数,但不能有位置参数。
  • 关键字参数名必须与依赖项的键名一致。
  • 依赖项必须通过 Provide 类声明。
  • 依赖项必须在处理函数的作用域内。

作用域:依赖项仅在其声明的上下文中可用。例如:

  • local_dependency 仅在声明它的路由处理函数中可用;
  • controller_dependency 仅在所属控制器的路由处理函数中可用;
  • router_dependency 仅在所属路由的路由处理函数中可用;
  • app_dependency 在所有路由处理函数中可用。

2.1. 带 yield 的依赖项(清理步骤)

依赖项可以是(异步)生成器函数,允许在处理函数返回后执行清理逻辑(如关闭连接)。
技术细节:清理步骤在处理函数返回后、响应发送前执行(针对 HTTP 请求)。

一个基本示例

from typing import Dict, Generator
from litestar import Litestar, get
from litestar.di import Provide

CONNECTION = {"open": False}

def generator_function() -> Generator[Dict[str, bool], None, None]:
    """打开连接,处理函数返回后关闭。"""
    CONNECTION["open"] = True
    yield CONNECTION  # 处理函数接收此处的返回值
    CONNECTION["open"] = False  # 清理步骤

@get("/", dependencies={"conn": Provide(generator_function)})
def index(conn: Dict[str, bool]) -> Dict[str, bool]:
    return conn

app = Litestar(route_handlers=[index])

如果你运行这段代码,你会看到在处理函数返回后,CONNECTION 已被重置:

from litestar.testing import TestClient
from dependencies import app, CONNECTION

with TestClient(app=app) as client:
    print(client.get("/").json())  # {"open": True}
    print(CONNECTION)  # {"open": False}

2.2. 异常处理

若处理函数中抛出异常,异常会在生成器的 yield 处触发,允许根据异常调整依赖行为(如数据库事务回滚)。

from collections.abc import Generator

from litestar import Litestar, get
from litestar.di import Provide

STATE = {"result": None, "connection": "closed"}


def generator_function() -> Generator[str, None, None]:
    """Set the connection state to open and close it after the handler returns.

    If an error occurs, set `result` to `"error"`, else set it to `"OK"`.
    """
    try:
        STATE["connection"] = "open"
        yield "hello"
        STATE["result"] = "OK"
    except ValueError:
        STATE["result"] = "error"
    finally:
        STATE["connection"] = "closed"


@get("/{name:str}", dependencies={"message": Provide(generator_function)})
def index(name: str, message: str) -> dict[str, str]:
    """If `name` is "John", return a message, otherwise raise an error."""
    if name == "John":
        return {name: message}
    raise ValueError()


app = Litestar(route_handlers=[index])
from litestar.testing import TestClient
from dependencies import STATE, app

with TestClient(app=app) as client:
    response = client.get("/John")
    print(response.json())  # {"John": "hello"}
    print(STATE)  # {"result": "OK", "connection": "closed"}

    response = client.get("/Peter")
    print(response.status_code)  # 500
    print(STATE)  # {"result": "error", "connection": "closed"}

最佳实践:始终使用 try/finally 包裹 yield,确保清理代码执行:

def generator_dependency():
    try:
        yield
    finally:
        ...  # cleanup code

3. 依赖项关键字参数

依赖项可接收 kwargs,因为它们与路由处理函数共享相同的参数解析机制,支持注入相同类型的数据(如路径参数、查询参数等)。

from litestar import Controller, patch
from litestar.di import Provide
from pydantic import BaseModel, UUID4

class User(BaseModel):
    id: UUID4
    name: str

async def retrieve_db_user(user_id: UUID4) -> User: ...  # 接收路径参数 user_id

class UserController(Controller):
    path = "/user"
    dependencies = {"user": Provide(retrieve_db_user)}  # 绑定依赖项

    @patch(path="/{user_id:uuid}")
    async def get_user(self, user: User) -> User: ...  # 注入 User 实例

4. 依赖项覆盖

通过相同键名在更低作用域(如路由处理函数)声明依赖项,可覆盖更高作用域(如控制器、路由、应用)的依赖项:

from litestar import Controller, get
from litestar.di import Provide


def bool_fn() -> bool: ...


def dict_fn() -> dict: ...


class MyController(Controller):
    # 控制器级依赖
    dependencies = {"some_dependency": Provide(dict_fn)}
    # 覆盖为布尔值
    @get(dependencies={"some_dependency": Provide(bool_fn)})
    def my_route_handler(self, some_dependency: bool) -> None: ...

较低范围的路由处理程序函数声明一个依赖项,其键与在较高范围的控制器上声明的键相同。因此,较低范围的依赖项会覆盖较高范围的依赖项。

5. Provide

Provide 类用于包装依赖项,确保其正确注入:

from litestar import get
from litestar.di import Provide
from random import randint

def my_dependency() -> int:
    return randint(1, 10)

@get("/some-path", dependencies={"my_dep": Provide(my_dependency)})
def my_handler(my_dep: int) -> None: ...

注意Provide.use_cache=True 时,依赖项返回值会被缓存(仅首次调用执行函数),使用时需谨慎。

6. 依赖中的依赖

依赖项可相互注入,规则与普通函数一致:

from litestar import Litestar, get
from litestar.di import Provide
from random import randint

def first_dependency() -> int:
    return randint(1, 10)

def second_dependency(injected_integer: int) -> bool:
    return injected_integer % 2 == 0  # 注入 first_dependency 的返回值

@get("/true-or-false")
def true_or_false_handler(injected_bool: bool) -> str:
    return "its true!" if injected_bool else "nope, its false..."

app = Litestar(
    route_handlers=[true_or_false_handler],
    dependencies={
        "injected_integer": Provide(first_dependency),
        "injected_bool": Provide(second_dependency),
    },
)

7. Dependency 函数

7.1. 依赖验证

默认情况下,Litestar 会验证注入值的类型。若类型不匹配,将抛出内部错误。可通过 Dependency(skip_validation=True) 跳过验证:

from typing import Annotated
from litestar.params import Dependency

def hello_world(injected: Annotated[int, Dependency(skip_validation=True)]) -> None: ...

7.2. 作为标记

从 OpenAPI 文档中排除带默认值的依赖项

声明 Dependency(default=3) 可避免参数被误判为查询参数:

def hello_world(optional_dependency: Annotated[int, Dependency(default=3)]) -> None: ...

未提供依赖时提前检测

标记 Dependency() 可在应用启动时检测必需依赖项缺失,避免运行时错误。