feat(typing): Accept extra context arguments in handler signatures

cq/_core/routing/di.py

@@ -3,7 +3,7 @@
 from abc import abstractmethod
 from collections.abc import Awaitable, Callable
 from contextlib import nullcontext
-from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
+from typing import TYPE_CHECKING, Any, Concatenate, Protocol, runtime_checkable
 
 from cq.middlewares.contextlib import AsyncContextManagerMiddleware
 
@@ -24,7 +24,7 @@ class DIAdapter(Protocol):
     __slots__ = ()
 
     @abstractmethod
-    def command_scope(self) -> Middleware[[Command], Any]:
+    def command_scope(self) -> Middleware[Concatenate[Command, ...], Any]:
         """
         Return a middleware that wraps each command dispatch.
 
@@ -88,7 +88,7 @@ def wire[T](self, tp: type[T]) -> Callable[..., Awaitable[T]]:
 class NoDI(DIAdapter):
     __slots__ = ()
 
-    def command_scope(self) -> Middleware[[Command], Any]:
+    def command_scope(self) -> Middleware[Concatenate[Command, ...], Any]:
         return AsyncContextManagerMiddleware(nullcontext())
 
     def lazy[T](self, tp: type[T], /) -> Callable[[], Awaitable[T]]:

cq/_core/routing/dispatchers/abc.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from collections.abc import Awaitable, Callable
-from typing import Protocol, Self, runtime_checkable
+from typing import Concatenate, Protocol, Self, runtime_checkable
 
 from cq._core.middleware import Middleware, MiddlewareGroup
 
@@ -20,18 +20,18 @@ async def dispatch(self, message: I, /) -> O:
 class BaseDispatcher[I, O](Dispatcher[I, O], ABC):
     __slots__ = ("__middleware_group",)
 
-    __middleware_group: MiddlewareGroup[[I], O]
+    __middleware_group: MiddlewareGroup[Concatenate[I, ...], O]
 
     def __init__(self) -> None:
         self.__middleware_group = MiddlewareGroup()
 
-    def add_middlewares(self, *middlewares: Middleware[[I], O]) -> Self:
+    def add_middlewares(self, *middlewares: Middleware[Concatenate[I, ...], O]) -> Self:
         self.__middleware_group.add(*middlewares)
         return self
 
     async def _invoke(
         self,
-        handler: Callable[[I], Awaitable[O]],
+        handler: Callable[Concatenate[I, ...], Awaitable[O]],
         message: I,
         /,
         fail_silently: bool = False,

cq/_core/routing/dispatchers/bus.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from collections.abc import Awaitable, Callable, Iterator
-from typing import Any, Protocol, Self, runtime_checkable
+from typing import Any, Concatenate, Protocol, Self, runtime_checkable
 
 import anyio
 from anyio.abc import TaskGroup
@@ -27,14 +27,14 @@ def add_listeners(self, *listeners: Listener[I]) -> Self:
         raise NotImplementedError
 
     @abstractmethod
-    def add_middlewares(self, *middlewares: Middleware[[I], O]) -> Self:
+    def add_middlewares(self, *middlewares: Middleware[Concatenate[I, ...], O]) -> Self:
         raise NotImplementedError
 
     @abstractmethod
     def subscribe(
         self,
         message_type: type[I],
-        factory: HandlerFactory[[I], O],
+        factory: HandlerFactory[Concatenate[I, ...], O],
         fail_silently: bool = ...,
     ) -> Self:
         raise NotImplementedError
@@ -58,13 +58,16 @@ def add_listeners(self, *listeners: Listener[I]) -> Self:
     def subscribe(
         self,
         message_type: type[I],
-        factory: HandlerFactory[[I], O],
+        factory: HandlerFactory[Concatenate[I, ...], O],
         fail_silently: bool = False,
     ) -> Self:
         self.__registry.subscribe(message_type, factory, fail_silently=fail_silently)
         return self
 
-    def _handlers_from(self, message_type: type[I]) -> Iterator[HandleFunction[[I], O]]:
+    def _handlers_from(
+        self,
+        message_type: type[I],
+    ) -> Iterator[HandleFunction[Concatenate[I, ...], O]]:
         return self.__registry.handlers_from(message_type)
 
     def _trigger_listeners(self, message: I, /, task_group: TaskGroup) -> None:

cq/_core/routing/handler.py

@@ -5,7 +5,15 @@
 from functools import partial
 from inspect import Parameter, isclass, unwrap
 from inspect import signature as inspect_signature
-from typing import TYPE_CHECKING, Any, Protocol, Self, overload, runtime_checkable
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Concatenate,
+    Protocol,
+    Self,
+    overload,
+    runtime_checkable,
+)
 
 from type_analyzer import MatchingTypesConfig, iter_matching_types, matching_types
 
@@ -55,23 +63,26 @@ def message_types(self) -> KeysView[type[I]]:
         raise NotImplementedError
 
     @abstractmethod
-    def handlers_from(self, message_type: type[I]) -> Iterator[HandleFunction[[I], O]]:
+    def handlers_from(
+        self,
+        message_type: type[I],
+    ) -> Iterator[HandleFunction[Concatenate[I, ...], O]]:
         raise NotImplementedError
 
     @abstractmethod
     def subscribe(
         self,
         message_type: type[I],
-        handler_factory: HandlerFactory[[I], O],
-        handler_type: HandlerType[[I], O] | None = ...,
+        handler_factory: HandlerFactory[Concatenate[I, ...], O],
+        handler_type: HandlerType[Concatenate[I, ...], O] | None = ...,
         fail_silently: bool = ...,
     ) -> Self:
         raise NotImplementedError
 
 
 @dataclass(repr=False, eq=False, frozen=True, slots=True)
 class MultipleHandlerRegistry[I, O](HandlerRegistry[I, O]):
-    __values: dict[type[I], list[HandleFunction[[I], O]]] = field(
+    __values: dict[type[I], list[HandleFunction[Concatenate[I, ...], O]]] = field(
         default_factory=partial(defaultdict, list),
         init=False,
     )
@@ -80,15 +91,18 @@ class MultipleHandlerRegistry[I, O](HandlerRegistry[I, O]):
     def message_types(self) -> KeysView[type[I]]:
         return self.__values.keys()
 
-    def handlers_from(self, message_type: type[I]) -> Iterator[HandleFunction[[I], O]]:
+    def handlers_from(
+        self,
+        message_type: type[I],
+    ) -> Iterator[HandleFunction[Concatenate[I, ...], O]]:
         for key_type in _iter_key_types(message_type):
             yield from self.__values.get(key_type, ())
 
     def subscribe(
         self,
         message_type: type[I],
-        handler_factory: HandlerFactory[[I], O],
-        handler_type: HandlerType[[I], O] | None = None,
+        handler_factory: HandlerFactory[Concatenate[I, ...], O],
+        handler_type: HandlerType[Concatenate[I, ...], O] | None = None,
         fail_silently: bool = False,
     ) -> Self:
         function = HandleFunction.create(handler_factory, handler_type, fail_silently)
@@ -101,7 +115,7 @@ def subscribe(
 
 @dataclass(repr=False, eq=False, frozen=True, slots=True)
 class SingleHandlerRegistry[I, O](HandlerRegistry[I, O]):
-    __values: dict[type[I], HandleFunction[[I], O]] = field(
+    __values: dict[type[I], HandleFunction[Concatenate[I, ...], O]] = field(
         default_factory=dict,
         init=False,
     )
@@ -110,7 +124,10 @@ class SingleHandlerRegistry[I, O](HandlerRegistry[I, O]):
     def message_types(self) -> KeysView[type[I]]:
         return self.__values.keys()
 
-    def handlers_from(self, message_type: type[I]) -> Iterator[HandleFunction[[I], O]]:
+    def handlers_from(
+        self,
+        message_type: type[I],
+    ) -> Iterator[HandleFunction[Concatenate[I, ...], O]]:
         for key_type in _iter_key_types(message_type):
             function = self.__values.get(key_type, None)
             if function is not None:
@@ -119,8 +136,8 @@ def handlers_from(self, message_type: type[I]) -> Iterator[HandleFunction[[I], O
     def subscribe(
         self,
         message_type: type[I],
-        handler_factory: HandlerFactory[[I], O],
-        handler_type: HandlerType[[I], O] | None = None,
+        handler_factory: HandlerFactory[Concatenate[I, ...], O],
+        handler_type: HandlerType[Concatenate[I, ...], O] | None = None,
         fail_silently: bool = False,
     ) -> Self:
         function = HandleFunction.create(handler_factory, handler_type, fail_silently)
@@ -195,12 +212,12 @@ def __call__[T](
 
     def __decorator(
         self,
-        wrapped: HandlerType[[I], O],
+        wrapped: HandlerType[Concatenate[I, ...], O],
         /,
         *,
         message_type: type[I] | None = None,
         fail_silently: bool = False,
-    ) -> HandlerType[[I], O]:
+    ) -> HandlerType[Concatenate[I, ...], O]:
         factory = self.di.wire(wrapped)
         message_type = message_type or _resolve_message_type(wrapped)
         self.registry.subscribe(message_type, factory, wrapped, fail_silently)
@@ -221,7 +238,9 @@ def _iter_key_types(message_type: Any) -> Iterator[Any]:
     return iter_matching_types(message_type, config)
 
 
-def _resolve_message_type[I, O](handler_type: HandlerType[[I], O]) -> type[I]:
+def _resolve_message_type[I, O](
+    handler_type: HandlerType[Concatenate[I, ...], O],
+) -> type[I]:
     fake_method = handler_type.handle.__get__(NotImplemented, handler_type)
     signature = inspect_signature(fake_method, eval_str=True)
 

cq/ext/injection.py

@@ -2,7 +2,7 @@
 from contextlib import AsyncExitStack
 from dataclasses import dataclass, field
 from enum import StrEnum
-from typing import Any
+from typing import Any, Concatenate
 
 from injection import Module, adefine_scope, mod
 from injection.exceptions import ScopeAlreadyDefinedError
@@ -24,7 +24,7 @@ class InjectionAdapter(DIAdapter):
     module: Module = field(default_factory=mod)
     threadsafe: bool | None = field(default=None)
 
-    def command_scope(self) -> Middleware[[Command], Any]:
+    def command_scope(self) -> Middleware[Concatenate[Command, ...], Any]:
         return InjectionScopeMiddleware(
             CQScope.COMMAND_DISPATCH,
             threadsafe=self.threadsafe,