From 541da2d7f8a418de5079dd3f1aede5e967fc8322 Mon Sep 17 00:00:00 2001 From: Pavel Kirilin Date: Sat, 14 Dec 2024 11:36:43 +0100 Subject: [PATCH] Added support for passing the FastaAPI instance. --- taskiq_fastapi/initializator.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/taskiq_fastapi/initializator.py b/taskiq_fastapi/initializator.py index 7ec8851..1c9e775 100644 --- a/taskiq_fastapi/initializator.py +++ b/taskiq_fastapi/initializator.py @@ -1,4 +1,4 @@ -from typing import Awaitable, Callable +from typing import Awaitable, Callable, Union from fastapi import FastAPI, Request from starlette.requests import HTTPConnection @@ -8,7 +8,7 @@ def startup_event_generator( broker: AsyncBroker, - app_path: str, + app_or_path: Union[str, FastAPI], ) -> Callable[[TaskiqState], Awaitable[None]]: """ Generate shutdown event. @@ -24,12 +24,16 @@ def startup_event_generator( async def startup(state: TaskiqState) -> None: if not broker.is_worker_process: return - app = import_object(app_path) + if isinstance(app_or_path, str): + app = import_object(app_or_path) + else: + app = app_or_path + if not isinstance(app, FastAPI): app = app() if not isinstance(app, FastAPI): - raise ValueError(f"'{app_path}' is not a FastAPI application.") + raise ValueError(f"'{app_or_path}' is not a FastAPI application.") state.fastapi_app = app await app.router.startup() @@ -62,7 +66,7 @@ async def shutdown(state: TaskiqState) -> None: return shutdown -def init(broker: AsyncBroker, app_path: str) -> None: +def init(broker: AsyncBroker, app_or_path: Union[str, FastAPI]) -> None: """ Add taskiq startup events. @@ -78,7 +82,7 @@ def init(broker: AsyncBroker, app_path: str) -> None: """ broker.add_event_handler( TaskiqEvents.WORKER_STARTUP, - startup_event_generator(broker, app_path), + startup_event_generator(broker, app_or_path), ) broker.add_event_handler(