feat: ✨ New handler registration logic

Author: remimd

cq/__init__.py

@@ -12,9 +12,9 @@
     QueryBus,
     command_handler,
     event_handler,
-    get_command_bus,
-    get_event_bus,
-    get_query_bus,
+    new_command_bus,
+    new_event_bus,
+    new_query_bus,
     query_handler,
 )
 from ._core.middleware import Middleware, MiddlewareResult
@@ -39,8 +39,8 @@
     "RelatedEvents",
     "command_handler",
     "event_handler",
-    "get_command_bus",
-    "get_event_bus",
-    "get_query_bus",
+    "new_command_bus",
+    "new_event_bus",
+    "new_query_bus",
     "query_handler",
 )

cq/_core/dispatcher/bus.py

@@ -1,32 +1,20 @@
 from abc import ABC, abstractmethod
-from collections import defaultdict
 from collections.abc import Awaitable, Callable
-from dataclasses import dataclass, field
-from inspect import getmro, isclass
-from types import GenericAlias
-from typing import Any, Protocol, Self, TypeAliasType, runtime_checkable
+from typing import Any, Protocol, Self, runtime_checkable
 
 import anyio
-import injection
+from anyio.abc import TaskGroup
 
 from cq._core.dispatcher.base import BaseDispatcher, Dispatcher
-
-type HandlerType[**P, T] = type[Handler[P, T]]
-type HandlerFactory[**P, T] = Callable[..., Awaitable[Handler[P, T]]]
+from cq._core.handler import (
+    HandlerFactory,
+    HandlerManager,
+    MultipleHandlerManager,
+    SingleHandlerManager,
+)
 
 type Listener[T] = Callable[[T], Awaitable[Any]]
 
-type BusType[I, O] = type[Bus[I, O]]
-
-
-@runtime_checkable
-class Handler[**P, T](Protocol):
-    __slots__ = ()
-
-    @abstractmethod
-    async def handle(self, *args: P.args, **kwargs: P.kwargs) -> T:
-        raise NotImplementedError
-
 
 @runtime_checkable
 class Bus[I, O](Dispatcher[I, O], Protocol):
@@ -41,27 +29,6 @@ def add_listeners(self, *listeners: Listener[I]) -> Self:
         raise NotImplementedError
 
 
-@dataclass(eq=False, frozen=True, slots=True)
-class SubscriberDecorator[I, O]:
-    bus_type: BusType[I, O] | TypeAliasType | GenericAlias
-    injection_module: injection.Module = field(default_factory=injection.mod)
-
-    def __call__(self, first_input_type: type[I], /, *input_types: type[I]) -> Any:
-        def decorator(wrapped: type[Handler[[I], O]]) -> type[Handler[[I], O]]:
-            if not isclass(wrapped) or not issubclass(wrapped, Handler):
-                raise TypeError(f"`{wrapped}` isn't a valid handler.")
-
-            bus = self.injection_module.find_instance(self.bus_type)
-            factory = self.injection_module.make_async_factory(wrapped)
-
-            for input_type in (first_input_type, *input_types):
-                bus.subscribe(input_type, factory)
-
-            return wrapped
-
-        return decorator
-
-
 class BaseBus[I, O](BaseDispatcher[I, O], Bus[I, O], ABC):
     __slots__ = ("__listeners",)
 
@@ -75,81 +42,47 @@ def add_listeners(self, *listeners: Listener[I]) -> Self:
         self.__listeners.extend(listeners)
         return self
 
-    async def _trigger_listeners(self, input_value: I, /) -> None:
-        listeners = self.__listeners
-
-        if not listeners:
-            return
-
-        async with anyio.create_task_group() as task_group:
-            for listener in listeners:
-                task_group.start_soon(listener, input_value)
-
-    @staticmethod
-    def _make_handle_function(
-        handler_factory: HandlerFactory[[I], O],
-    ) -> Callable[[I], Awaitable[O]]:
-        async def handle(input_value: I) -> O:
-            handler = await handler_factory()
-            return await handler.handle(input_value)
-
-        return handle
+    def _trigger_listeners(self, input_value: I, /, task_group: TaskGroup) -> None:
+        for listener in self.__listeners:
+            task_group.start_soon(listener, input_value)
 
 
 class SimpleBus[I, O](BaseBus[I, O]):
-    __slots__ = ("__handlers",)
+    __slots__ = ("__manager",)
 
-    __handlers: dict[type[I], HandlerFactory[[I], O]]
+    __manager: HandlerManager[I, O]
 
-    def __init__(self) -> None:
+    def __init__(self, manager: HandlerManager[I, O] | None = None) -> None:
         super().__init__()
-        self.__handlers = {}
+        self.__manager = manager or SingleHandlerManager()
 
     async def dispatch(self, input_value: I, /) -> O:
-        await self._trigger_listeners(input_value)
-
-        for input_type in getmro(type(input_value)):
-            if handler_factory := self.__handlers.get(input_type):
-                break
+        async with anyio.create_task_group() as task_group:
+            self._trigger_listeners(input_value, task_group)
 
-        else:
-            return NotImplemented
+        for handler in self.__manager.handlers_from(type(input_value)):
+            return await self._invoke_with_middlewares(handler, input_value)
 
-        handler = self._make_handle_function(handler_factory)
-        return await self._invoke_with_middlewares(handler, input_value)
+        return NotImplemented
 
     def subscribe(self, input_type: type[I], factory: HandlerFactory[[I], O]) -> Self:
-        if input_type in self.__handlers:
-            raise RuntimeError(
-                f"A handler is already registered for the input type: `{input_type}`."
-            )
-
-        self.__handlers[input_type] = factory
+        self.__manager.subscribe(input_type, factory)
         return self
 
 
 class TaskBus[I](BaseBus[I, None]):
-    __slots__ = ("__handlers",)
+    __slots__ = ("__manager",)
 
-    __handlers: dict[type[I], list[HandlerFactory[[I], None]]]
+    __manager: HandlerManager[I, None]
 
-    def __init__(self) -> None:
+    def __init__(self, manager: HandlerManager[I, None] | None = None) -> None:
         super().__init__()
-        self.__handlers = defaultdict(list)
+        self.__manager = manager or MultipleHandlerManager()
 
     async def dispatch(self, input_value: I, /) -> None:
-        await self._trigger_listeners(input_value)
-
-        for input_type in getmro(type(input_value)):
-            if handler_factories := self.__handlers.get(input_type):
-                break
-
-        else:
-            return
-
         async with anyio.create_task_group() as task_group:
-            for handler_factory in handler_factories:
-                handler = self._make_handle_function(handler_factory)
+            self._trigger_listeners(input_value, task_group)
+            for handler in self.__manager.handlers_from(type(input_value)):
                 task_group.start_soon(
                     self._invoke_with_middlewares,
                     handler,
@@ -161,5 +94,5 @@ def subscribe(
         input_type: type[I],
         factory: HandlerFactory[[I], None],
     ) -> Self:
-        self.__handlers[input_type].append(factory)
+        self.__manager.subscribe(input_type, factory)
         return self

cq/_core/handler.py

@@ -0,0 +1,111 @@
+from abc import abstractmethod
+from collections import defaultdict
+from collections.abc import Awaitable, Callable, Iterator
+from dataclasses import dataclass, field
+from functools import partial
+from inspect import getmro, isclass
+from typing import Any, Protocol, Self, runtime_checkable
+
+import injection
+
+type HandlerType[**P, T] = type[Handler[P, T]]
+type HandlerFactory[**P, T] = Callable[..., Awaitable[Handler[P, T]]]
+
+
+@runtime_checkable
+class Handler[**P, T](Protocol):
+    __slots__ = ()
+
+    @abstractmethod
+    async def handle(self, *args: P.args, **kwargs: P.kwargs) -> T:
+        raise NotImplementedError
+
+
+@runtime_checkable
+class HandlerManager[I, O](Protocol):
+    __slots__ = ()
+
+    @abstractmethod
+    def handlers_from(
+        self,
+        input_type: type[I],
+    ) -> Iterator[Callable[[I], Awaitable[O]]]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def subscribe(self, input_type: type[I], factory: HandlerFactory[[I], O]) -> Self:
+        raise NotImplementedError
+
+
+@dataclass(eq=False, frozen=True, slots=True)
+class MultipleHandlerManager[I, O](HandlerManager[I, O]):
+    __factories: dict[type[I], list[HandlerFactory[[I], O]]] = field(
+        default_factory=partial(defaultdict, list),
+        init=False,
+    )
+
+    def handlers_from(
+        self,
+        input_type: type[I],
+    ) -> Iterator[Callable[[I], Awaitable[O]]]:
+        for it in getmro(input_type):
+            for factory in self.__factories.get(it, ()):
+                yield _make_handle_function(factory)
+
+    def subscribe(self, input_type: type[I], factory: HandlerFactory[[I], O]) -> Self:
+        self.__factories[input_type].append(factory)
+        return self
+
+
+@dataclass(eq=False, frozen=True, slots=True)
+class SingleHandlerManager[I, O](HandlerManager[I, O]):
+    __factories: dict[type[I], HandlerFactory[[I], O]] = field(
+        default_factory=dict,
+        init=False,
+    )
+
+    def handlers_from(
+        self,
+        input_type: type[I],
+    ) -> Iterator[Callable[[I], Awaitable[O]]]:
+        for it in getmro(input_type):
+            factory = self.__factories.get(it, None)
+            if factory is not None:
+                yield _make_handle_function(factory)
+
+    def subscribe(self, input_type: type[I], factory: HandlerFactory[[I], O]) -> Self:
+        if input_type in self.__factories:
+            raise RuntimeError(
+                f"A handler is already registered for the input type: `{input_type}`."
+            )
+
+        self.__factories[input_type] = factory
+        return self
+
+
+@dataclass(eq=False, frozen=True, slots=True)
+class HandlerDecorator[I, O]:
+    manager: HandlerManager[I, O]
+    injection_module: injection.Module = field(default_factory=injection.mod)
+
+    def __call__(self, input_type: type[I], /) -> Any:
+        def decorator(wrapped: type[Handler[[I], O]]) -> type[Handler[[I], O]]:
+            if not isclass(wrapped) or not issubclass(wrapped, Handler):
+                raise TypeError(f"`{wrapped}` isn't a valid handler.")
+
+            factory = self.injection_module.make_async_factory(wrapped)
+            self.manager.subscribe(input_type, factory)
+            return wrapped
+
+        return decorator
+
+
+def _make_handle_function[I, O](
+    handler_factory: HandlerFactory[[I], O],
+) -> Callable[[I], Awaitable[O]]:
+    return partial(__handle, factory=handler_factory)
+
+
+async def __handle[I, O](input_value: I, factory: HandlerFactory[[I], O]) -> O:
+    handler = await factory()
+    return await handler.handle(input_value)

cq/_core/message.py

@@ -3,8 +3,13 @@
 
 import injection
 
-from cq._core.dispatcher.bus import Bus, SimpleBus, SubscriberDecorator, TaskBus
+from cq._core.dispatcher.bus import Bus, SimpleBus, TaskBus
 from cq._core.dto import DTO
+from cq._core.handler import (
+    HandlerDecorator,
+    MultipleHandlerManager,
+    SingleHandlerManager,
+)
 from cq._core.scope import CQScope
 from cq.middlewares.scope import InjectionScopeMiddleware
 
@@ -32,42 +37,42 @@ class Query(Message, ABC):
 AnyCommandBus = CommandBus[Any]
 
 
-command_handler: SubscriberDecorator[Command, Any] = SubscriberDecorator(CommandBus)
-event_handler: SubscriberDecorator[Event, None] = SubscriberDecorator(EventBus)
-query_handler: SubscriberDecorator[Query, Any] = SubscriberDecorator(QueryBus)
-
-
-@injection.inject
-def get_command_bus[T](bus: CommandBus[T] = NotImplemented, /) -> CommandBus[T]:
-    return bus
+command_handler: HandlerDecorator[Command, Any] = HandlerDecorator(
+    SingleHandlerManager(),
+)
+event_handler: HandlerDecorator[Event, None] = HandlerDecorator(
+    MultipleHandlerManager(),
+)
+query_handler: HandlerDecorator[Query, Any] = HandlerDecorator(
+    SingleHandlerManager(),
+)
 
 
+@injection.singleton(
+    on=CommandBus,
+    ignore_type_hint=True,  # type: ignore[call-arg]
+    inject=False,
+    mode="fallback",
+)
 def new_command_bus[T]() -> CommandBus[T]:
-    bus: CommandBus[T] = SimpleBus()
-    bus.add_middlewares(
+    return SimpleBus(command_handler.manager).add_middlewares(
         InjectionScopeMiddleware(CQScope.ON_COMMAND),
     )
-    return bus
-
-
-@injection.inject
-def get_event_bus(bus: EventBus = NotImplemented, /) -> EventBus:
-    return bus
 
 
+@injection.singleton(
+    inject=False,
+    mode="fallback",
+)
 def new_event_bus() -> EventBus:
-    return TaskBus()
-
-
-@injection.inject
-def get_query_bus[T](bus: QueryBus[T] = NotImplemented, /) -> QueryBus[T]:
-    return bus
+    return TaskBus(event_handler.manager)
 
 
+@injection.singleton(
+    on=QueryBus,
+    ignore_type_hint=True,  # type: ignore[call-arg]
+    inject=False,
+    mode="fallback",
+)
 def new_query_bus[T]() -> QueryBus[T]:
-    return SimpleBus()
-
-
-injection.set_constant(new_command_bus(), CommandBus, alias=True)
-injection.set_constant(new_event_bus(), EventBus, alias=True)
-injection.set_constant(new_query_bus(), QueryBus, alias=True)
+    return SimpleBus(query_handler.manager)