feat: ✨ Introduce ContextPipeline

remimd · web-flow · commit affa3c367e18 · 2025-10-21T00:54:07.000+02:00
diff --git a/cq/__init__.py b/cq/__init__.py
@@ -1,6 +1,7 @@
 from ._core.dispatcher.base import DeferredDispatcher, Dispatcher
 from ._core.dispatcher.bus import Bus
-from ._core.dispatcher.pipe import Pipe
+from ._core.dispatcher.lazy import LazyDispatcher
+from ._core.dispatcher.pipe import ContextPipeline, Pipe
 from ._core.message import (
     AnyCommandBus,
     Command,
@@ -17,6 +18,7 @@
     query_handler,
 )
 from ._core.middleware import Middleware, MiddlewareResult
+from ._core.pipetools import ContextCommandPipeline
 from ._core.related_events import RelatedEvents
 from ._core.scope import CQScope
 
@@ -26,10 +28,13 @@
     "CQScope",
     "Command",
     "CommandBus",
+    "ContextCommandPipeline",
+    "ContextPipeline",
     "DeferredDispatcher",
     "Dispatcher",
     "Event",
     "EventBus",
+    "LazyDispatcher",
     "Middleware",
     "MiddlewareResult",
     "Pipe",
diff --git a/cq/_core/dispatcher/lazy.py b/cq/_core/dispatcher/lazy.py
@@ -0,0 +1,28 @@
+from collections.abc import Awaitable
+from types import GenericAlias
+from typing import TypeAliasType
+
+import injection
+
+from cq._core.dispatcher.base import Dispatcher
+
+
+class LazyDispatcher[I, O](Dispatcher[I, O]):
+    __slots__ = ("__value",)
+
+    __value: Awaitable[Dispatcher[I, O]]
+
+    def __init__(
+        self,
+        dispatcher_type: type[Dispatcher[I, O]] | TypeAliasType | GenericAlias,
+        /,
+        *,
+        injection_module: injection.Module | None = None,
+        threadsafe: bool | None = None,
+    ) -> None:
+        module = injection_module or injection.mod()
+        self.__value = module.aget_lazy_instance(dispatcher_type, threadsafe=threadsafe)
+
+    async def dispatch(self, input_value: I, /) -> O:
+        dispatcher = await self.__value
+        return await dispatcher.dispatch(input_value)
diff --git a/cq/_core/dispatcher/pipe.py b/cq/_core/dispatcher/pipe.py
@@ -1,12 +1,22 @@
+from collections import deque
 from collections.abc import Awaitable, Callable
 from dataclasses import dataclass, field
-from typing import Any, Self, overload
+from typing import Any, Protocol, Self, overload
 
 from cq._core.dispatcher.base import BaseDispatcher, Dispatcher
+from cq._core.middleware import Middleware
 
 type PipeConverter[I, O] = Callable[[O], Awaitable[I]]
 
 
+class PipeConverterMethod[I, O](Protocol):
+    def __get__(
+        self,
+        instance: object,
+        owner: type | None = ...,
+    ) -> PipeConverter[I, O]: ...
+
+
 @dataclass(repr=False, eq=False, frozen=True, slots=True)
 class PipeStep[I, O]:
     converter: PipeConverter[I, O]
@@ -77,6 +87,110 @@ async def __execute(self, input_value: I) -> O:
         for step in self.__steps:
             output_value = await dispatcher.dispatch(input_value)
             input_value = await step.converter(output_value)
+
+            if input_value is None:
+                return NotImplemented
+
             dispatcher = step.dispatcher or self.__dispatcher
 
         return await dispatcher.dispatch(input_value)
+
+
+@dataclass(repr=False, eq=False, frozen=True, slots=True)
+class ContextPipelineStep[I, O]:
+    converter: PipeConverterMethod[I, O]
+    dispatcher: Dispatcher[I, Any] | None = field(default=None)
+
+
+class ContextPipeline[I]:
+    __slots__ = ("__dispatcher", "__middlewares", "__steps")
+
+    __dispatcher: Dispatcher[Any, Any]
+    __middlewares: deque[Middleware[Any, Any]]
+    __steps: list[ContextPipelineStep[Any, Any]]
+
+    def __init__(self, dispatcher: Dispatcher[Any, Any]) -> None:
+        self.__dispatcher = dispatcher
+        self.__middlewares = deque()
+        self.__steps = []
+
+    @overload
+    def __get__[O](
+        self,
+        instance: O,
+        owner: type[O] | None = ...,
+    ) -> Dispatcher[I, O]: ...
+
+    @overload
+    def __get__(self, instance: None = ..., owner: type | None = ...) -> Self: ...
+
+    def __get__[O](
+        self,
+        instance: O | None = None,
+        owner: type[O] | None = None,
+    ) -> Self | Dispatcher[I, O]:
+        if instance is None:
+            return self
+
+        pipeline = self.__new_pipeline(instance, owner)
+        return BoundContextPipeline(instance, pipeline)
+
+    def add_middlewares(self, *middlewares: Middleware[[I], Any]) -> Self:
+        self.__middlewares.extendleft(reversed(middlewares))
+        return self
+
+    @overload
+    def step[T](
+        self,
+        wrapped: PipeConverterMethod[T, Any],
+        /,
+        *,
+        dispatcher: Dispatcher[T, Any] | None = ...,
+    ) -> PipeConverterMethod[T, Any]: ...
+
+    @overload
+    def step[T](
+        self,
+        wrapped: None = ...,
+        /,
+        *,
+        dispatcher: Dispatcher[T, Any] | None = ...,
+    ) -> Callable[[PipeConverterMethod[T, Any]], PipeConverterMethod[T, Any]]: ...
+
+    def step[T](
+        self,
+        wrapped: PipeConverterMethod[T, Any] | None = None,
+        /,
+        *,
+        dispatcher: Dispatcher[T, Any] | None = None,
+    ) -> Any:
+        def decorator(wp: PipeConverterMethod[T, Any]) -> PipeConverterMethod[T, Any]:
+            step = ContextPipelineStep(wp, dispatcher)
+            self.__steps.append(step)
+            return wp
+
+        return decorator(wrapped) if wrapped else decorator
+
+    def __new_pipeline[T](
+        self,
+        context: T,
+        context_type: type[T] | None,
+    ) -> Pipe[I, Any]:
+        pipeline: Pipe[I, Any] = Pipe(self.__dispatcher)
+        pipeline.add_middlewares(*self.__middlewares)
+
+        for step in self.__steps:
+            converter = step.converter.__get__(context, context_type)
+            pipeline.step(converter, dispatcher=step.dispatcher)
+
+        return pipeline
+
+
+@dataclass(repr=False, eq=False, frozen=True, slots=True)
+class BoundContextPipeline[I, O](Dispatcher[I, O]):
+    context: O
+    pipeline: Pipe[I, Any]
+
+    async def dispatch(self, input_value: I, /) -> O:
+        await self.pipeline.dispatch(input_value)
+        return self.context
diff --git a/cq/_core/pipetools.py b/cq/_core/pipetools.py
@@ -0,0 +1,31 @@
+import injection
+
+from cq._core.dispatcher.lazy import LazyDispatcher
+from cq._core.dispatcher.pipe import ContextPipeline
+from cq._core.message import AnyCommandBus, Command
+from cq._core.scope import CQScope
+from cq.middlewares.scope import InjectionScopeMiddleware
+
+
+class ContextCommandPipeline[I: Command](ContextPipeline[I]):
+    __slots__ = ()
+
+    def __init__(
+        self,
+        /,
+        *,
+        injection_module: injection.Module | None = None,
+        threadsafe: bool | None = None,
+    ) -> None:
+        dispatcher = LazyDispatcher(
+            AnyCommandBus,
+            injection_module=injection_module,
+            threadsafe=threadsafe,
+        )
+        super().__init__(dispatcher)
+        transaction_scope_middleware = InjectionScopeMiddleware(
+            CQScope.TRANSACTION,
+            exist_ok=True,
+            threadsafe=threadsafe,
+        )
+        self.add_middlewares(transaction_scope_middleware)
diff --git a/documentation/pipeline.md b/documentation/pipeline.md
@@ -2,6 +2,10 @@
 
 Pipelines are designed to execute several commands one after the other, while encapsulating them in middleware.
 
+## Pipe
+
+A simple pipeline implementation.
+
 Example:
 
 ```python
@@ -12,14 +16,54 @@ async def pipeline_example(command_bus: AnyCommandBus) -> None:
     pipeline.add_middlewares(...)  # You can add middleware to encapsulate the pipeline in a transaction, for example.
 
     @pipeline.step
-    async def converter_1(output_value: FirstResult) -> SecondCommand:
+    async def _(result: FirstResult) -> SecondCommand:
         """ Transform the return value into a new command. """
 
     @pipeline.step
-    async def converter_2(output_value: SecondResult) -> ThirdCommand:
+    async def _(result: SecondResult) -> ThirdCommand:
         """ Transform the return value into a new command. """
 
     command = FirstCommand(...)
     output_value = await pipeline.dispatch(command)
     # ...
 ```
+
+## ContextPipeline
+
+A limitation of `Pipe` is that it isn't possible to have values that go through the steps. Each converter in a `Pipe` 
+receives only the output of the previous one, but the intermediate state, any contextual information, is lost between 
+steps.
+
+This makes it impossible to accumulate data or share context between converters without resorting to global variables.
+
+To solve this limitation, `ContextPipeline` introduces a contextual layer: each stage operates within the same instance,
+allowing stateful processing and side effects to persist across the entire pipeline execution.
+
+Example:
+
+```python
+from cq import ContextCommandPipeline, ContextPipeline
+
+class ContextExample:
+    user_id: int
+
+    pipeline: ContextPipeline[FirstCommand] = ContextCommandPipeline()
+
+    @pipeline.step
+    async def _(self, result: FirstResult) -> SecondCommand:
+        self.user_id = result.user_id  # set user_id in context
+        return SecondCommand(...)
+
+    @pipeline.step
+    async def _(self, result: SecondResult) -> ThirdCommand:
+        """ Transform the return value into a new command. """
+        
+    @pipeline.step
+    async def _(self, result: ThirdResult) -> None:
+        """ The last step is optional, but if you need it, you must return `None`. """
+
+async def how_to_dispatch() -> None:
+    command = FirstCommand(...)
+    context = await ContextExample().pipeline.dispatch(command)
+    # ...
+```
diff --git a/tests/test_context_command_pipeline.py b/tests/test_context_command_pipeline.py
@@ -0,0 +1,65 @@
+from cq import ContextCommandPipeline, ContextPipeline, command_handler
+from tests.helpers.history import HistoryMiddleware
+
+
+class TestContextCommandPipeline:
+    async def test_dispatch_with_success_return_any(
+        self,
+        history: HistoryMiddleware,
+    ) -> None:
+        class Command1: ...
+
+        class Command2: ...
+
+        class Command3: ...
+
+        class Foo: ...
+
+        class Bar: ...
+
+        class Baz: ...
+
+        @command_handler
+        class CommandHandler1:
+            async def handle(self, command: Command1) -> Foo:
+                return Foo()
+
+        @command_handler
+        class CommandHandler2:
+            async def handle(self, command: Command2) -> Bar:
+                return Bar()
+
+        @command_handler
+        class CommandHandler3:
+            async def handle(self, command: Command3) -> Baz:
+                return Baz()
+
+        class Context:
+            foo: Foo
+            bar: Bar
+            baz: Baz
+
+            pipeline: ContextPipeline[Command1] = ContextCommandPipeline()
+
+            @pipeline.step
+            async def _(self, foo: Foo) -> Command2:
+                self.foo = foo
+                return Command2()
+
+            @pipeline.step
+            async def _(self, bar: Bar) -> Command3:
+                self.bar = bar
+                return Command3()
+
+            @pipeline.step
+            async def _(self, baz: Baz) -> None:
+                self.baz = baz
+
+        cmd = Command1()
+        ctx = await Context().pipeline.dispatch(cmd)
+
+        assert isinstance(ctx, Context)
+        assert isinstance(ctx.foo, Foo)
+        assert isinstance(ctx.bar, Bar)
+        assert isinstance(ctx.baz, Baz)
+        assert len(history.records) == 3
diff --git a/uv.lock b/uv.lock
