diff --git a/README.md b/README.md index 94d1dce..c1242a3 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,7 @@ async def greet( await message.answer(f"{settings.greeting}, {message.from_user.first_name}") ``` -Pass `auto_inject=True` to `setup_di` to wrap every handler already registered on the dispatcher, so individual handlers don't need `@inject`. Register handlers before startup for this to take effect. The current `aiogram.types.Update` and the concrete event it carries (`Message`, `CallbackQuery`, and so on) are resolvable within DI via the pre-built `aiogram_update_provider` / `aiogram_event_provider` context providers. [aiogram-dialog](https://github.com/Tishka17/aiogram_dialog) getters and callbacks are supported via `modern_di_aiogram.dialog`; see the docs. +Pass `auto_inject=True` to `setup_di` to wrap every handler already registered on the dispatcher, so individual handlers don't need `@inject`. Register handlers before startup for this to take effect. Sync handlers work with both `@inject` and `auto_inject`: aiogram runs them in a worker thread as usual, and their `FromDI` parameters resolve from the same per-update child container. The current `aiogram.types.Update` and the concrete event it carries (`Message`, `CallbackQuery`, and so on) are resolvable within DI via the pre-built `aiogram_update_provider` / `aiogram_event_provider` context providers. [aiogram-dialog](https://github.com/Tishka17/aiogram_dialog) getters and callbacks are supported via `modern_di_aiogram.dialog`; see the docs. ## API @@ -74,7 +74,7 @@ Pass `auto_inject=True` to `setup_di` to wrap every handler already registered o |---|---| | `setup_di(dispatcher, container, *, auto_inject=False)` | Stores the container on the dispatcher, registers the update/event providers, wires `dispatcher.startup`/`dispatcher.shutdown` to open/close it, and installs the per-update middleware. With `auto_inject=True`, also wraps every handler already registered at startup | | `FromDI(dependency)` | Inert marker (used with `@inject`) that resolves a provider or type from the per-update child container | -| `inject(handler)` | Decorator for an aiogram handler; resolves its `FromDI`-annotated parameters. Not needed when `setup_di(..., auto_inject=True)` is used. Raises `RuntimeError` naming `setup_di` when an update reaches it without the middleware installed | +| `inject(handler)` | Decorator for an aiogram handler, sync or async; resolves its `FromDI`-annotated parameters and returns a wrapper of the same kind. Not needed when `setup_di(..., auto_inject=True)` is used. Raises `RuntimeError` naming `setup_di` when an update reaches it without the middleware installed | | `fetch_di_container(dispatcher)` | Returns the root `Container` stored on the dispatcher | | `aiogram_update_provider` | `ContextProvider` for the current `aiogram.types.Update` (`REQUEST` scope) | | `aiogram_event_provider` | `ContextProvider` for the current `aiogram.types.TelegramObject` (`REQUEST` scope): the concrete event unwrapped from the `Update` | diff --git a/modern_di_aiogram/main.py b/modern_di_aiogram/main.py index 1a0b11b..929f9a9 100644 --- a/modern_di_aiogram/main.py +++ b/modern_di_aiogram/main.py @@ -71,7 +71,7 @@ def fetch_di_container(dispatcher: Dispatcher) -> Container: FromDI = integrations.from_di -def inject(func: typing.Callable[..., typing.Awaitable[T]]) -> typing.Callable[..., typing.Awaitable[T]]: +def inject(func: typing.Callable[..., T]) -> typing.Callable[..., T]: di_params = integrations.parse_markers(func) if not di_params: integrations.mark_injected(func) @@ -85,21 +85,36 @@ def inject(func: typing.Callable[..., typing.Awaitable[T]]) -> typing.Callable[. inspect.Parameter(_CHILD_CONTAINER_KEY, kind=inspect.Parameter.KEYWORD_ONLY, annotation=Container), ) - async def wrapper(*args: typing.Any, **kwargs: typing.Any) -> T: # noqa: ANN401 + def resolve(kwargs: dict[str, typing.Any]) -> dict[str, typing.Any]: container = _fetch_child_container(kwargs) if container_param_injected: del kwargs[_CHILD_CONTAINER_KEY] - resolved = integrations.resolve_markers(container, di_params) - return await func(*args, **kwargs, **resolved) + return integrations.resolve_markers(container, di_params) + + wrapper: typing.Callable[..., typing.Any] + if inspect.iscoroutinefunction(func): + + async def async_wrapper(*args: typing.Any, **kwargs: typing.Any) -> typing.Any: # noqa: ANN401 + resolved = resolve(kwargs) + return await func(*args, **kwargs, **resolved) + + wrapper = async_wrapper + else: + + def sync_wrapper(*args: typing.Any, **kwargs: typing.Any) -> T: # noqa: ANN401 + resolved = resolve(kwargs) + return func(*args, **kwargs, **resolved) + + wrapper = sync_wrapper # NOT functools.wraps: aiogram unwraps __wrapped__ and would defeat __signature__. wrapper.__name__ = func.__name__ # ty: ignore[unresolved-attribute] wrapper.__qualname__ = func.__qualname__ # ty: ignore[unresolved-attribute] wrapper.__doc__ = func.__doc__ wrapper.__module__ = func.__module__ - wrapper.__signature__ = original_signature.replace(parameters=visible_params) # ty: ignore[unresolved-attribute] + wrapper.__signature__ = original_signature.replace(parameters=visible_params) # ty: ignore[invalid-assignment] integrations.mark_injected(wrapper) - return wrapper + return typing.cast(typing.Callable[..., T], wrapper) def _inject_router(router: aiogram.Router) -> None: @@ -115,4 +130,6 @@ def _inject_router(router: aiogram.Router) -> None: flags=handler.flags, ) handler.callback = wrapped.callback + handler.awaitable = wrapped.awaitable handler.params = wrapped.params + handler.varkw = wrapped.varkw diff --git a/tests/test_auto_inject.py b/tests/test_auto_inject.py index 6869f1a..b60a20c 100644 --- a/tests/test_auto_inject.py +++ b/tests/test_auto_inject.py @@ -1,4 +1,6 @@ +import inspect import typing +import warnings from aiogram import Bot, Dispatcher from aiogram.types import Message @@ -46,3 +48,48 @@ async def handler( await dispatcher.emit_shutdown() assert seen == {"app": "original"} + + +async def test_auto_inject_runs_sync_handler(bot: Bot) -> None: + dispatcher = Dispatcher() + setup_di(dispatcher, Container(groups=[Dependencies]), auto_inject=True) + seen: dict[str, typing.Any] = {} + + @dispatcher.message() + def handler( + message: Message, # noqa: ARG001 + app_instance: typing.Annotated[SimpleCreator, FromDI(SimpleCreator)], + ) -> None: + seen["dep1"] = app_instance.dep1 + + await dispatcher.emit_startup() + with warnings.catch_warnings(): + warnings.simplefilter("error") + await dispatcher.feed_update(bot, make_message_update()) + await dispatcher.emit_shutdown() + + assert seen == {"dep1": "original"} + + +async def test_auto_inject_keeps_handler_object_in_sync_with_callback() -> None: + dispatcher = Dispatcher() + setup_di(dispatcher, Container(groups=[Dependencies]), auto_inject=True) + + @dispatcher.message() + def sync_handler( + message: Message, + app_instance: typing.Annotated[SimpleCreator, FromDI(SimpleCreator)], + ) -> None: ... + + @dispatcher.message() + async def async_handler( + message: Message, + app_instance: typing.Annotated[SimpleCreator, FromDI(SimpleCreator)], + ) -> None: ... + + await dispatcher.emit_startup() + await dispatcher.emit_shutdown() + + assert [h.awaitable for h in dispatcher.message.handlers] == [ + inspect.iscoroutinefunction(h.callback) for h in dispatcher.message.handlers + ] diff --git a/tests/test_inject.py b/tests/test_inject.py index 4187ee8..8548c5d 100644 --- a/tests/test_inject.py +++ b/tests/test_inject.py @@ -1,5 +1,7 @@ import contextlib +import inspect import typing +import warnings import pytest from aiogram import Bot, Dispatcher @@ -113,3 +115,52 @@ async def handler( with pytest.raises(RuntimeError, match=r"setup_di\(dispatcher, container\)"): await dispatcher.feed_update(bot, make_message_update()) + + +async def test_inject_runs_sync_handler_once(dispatcher: Dispatcher, bot: Bot) -> None: + calls: list[str] = [] + + @dispatcher.message() + @inject + def handler( + message: Message, # noqa: ARG001 + app_instance: typing.Annotated[SimpleCreator, FromDI(SimpleCreator)], + ) -> None: + calls.append(app_instance.dep1) + + await dispatcher.emit_startup() + with warnings.catch_warnings(): + warnings.simplefilter("error") + await dispatcher.feed_update(bot, make_message_update()) + await dispatcher.emit_shutdown() + + assert calls == ["original"] + + +async def test_inject_sync_handler_forwards_declared_container_param(dispatcher: Dispatcher, bot: Bot) -> None: + seen: dict[str, typing.Any] = {} + + @dispatcher.message() + @inject + def handler( + message: Message, # noqa: ARG001 + modern_di_container: Container, + app_instance: typing.Annotated[SimpleCreator, FromDI(SimpleCreator)], + ) -> None: + seen["got_container"] = isinstance(modern_di_container, Container) + seen["app"] = app_instance.dep1 + + await dispatcher.emit_startup() + await dispatcher.feed_update(bot, make_message_update()) + await dispatcher.emit_shutdown() + + assert seen == {"got_container": True, "app": "original"} + + +def test_inject_wrapper_is_async_only_for_async_handlers() -> None: + def sync_handler(app: typing.Annotated[SimpleCreator, FromDI(SimpleCreator)]) -> None: ... + + async def async_handler(app: typing.Annotated[SimpleCreator, FromDI(SimpleCreator)]) -> None: ... + + assert not inspect.iscoroutinefunction(inject(sync_handler)) + assert inspect.iscoroutinefunction(inject(async_handler))