feat: #2693 Add ContextRepo composition

docs/docs/en/getting-started/context.md

@@ -80,16 +80,66 @@ Afterward, you can access your `secret` field in the usual way:
 
 In this case, the field becomes a global context field: it does not depend on the current message handler (unlike `message`)
 
-!!! tip
-    Alternatively you can setup global context objects in `FastStream` object constructor:
+Alternatively you can setup global context objects in `FastStream` object constructor:
+
+```python
+from faststream import FastStream
+from faststream.context import ContextRepo
+
+context = ContextRepo({"secret_str": "my-perfect-secret"})
+app = FastStream(context=context)
+```
+
+And in brokers:
+
+=== "AIOKafka"
+
+    ```python
+    from faststream.context import ContextRepo
+    from faststream.kafka import KafkaBroker
+
+    context = ContextRepo({"secret_str": "my-perfect-secret"})
+    broker = KafkaBroker(context=context)
+    ```
+
+=== "Confluent"
+
+    ```python
+    from faststream.context import ContextRepo
+    from faststream.confluent import KafkaBroker
+
+    context = ContextRepo({"secret_str": "my-perfect-secret"})
+    broker = KafkaBroker(context=context)
+    ```
+
+=== "RabbitMQ"
 
     ```python
-    from faststream import FastStream
     from faststream.context import ContextRepo
+    from faststream.rabbit import RabbitBroker
 
-    app = FastStream(context=ContextRepo({
-        "secret_str": "my-perfect-secret"
-    }))
+    context = ContextRepo({"secret_str": "my-perfect-secret"})
+    broker = RabbitBroker(context=context)
+    ```
+
+=== "NATS"
+
+    ```python
+    from faststream.context import ContextRepo
+    from faststream.nats import NatsBroker
+
+    context = ContextRepo({"secret_str": "my-perfect-secret"})
+    broker = NatsBroker(context=context)
+    ```
+
+=== "Redis"
+
+    ```python
+    from faststream.context import ContextRepo
+    from faststream.redis import RedisBroker
+
+    context = ContextRepo({"secret_str": "my-perfect-secret"})
+    broker = RedisBroker(context=context)
     ```
 
 To remove a field from the context use the `reset_global` method:
@@ -98,6 +148,22 @@ To remove a field from the context use the `reset_global` method:
 context.reset_global("my_key")
 ```
 
+!!! tip
+    It is important to keep in mind that the broker context takes precedence over the FastStream context.
+
+    ```python
+    from faststream import FastStream
+    from faststream.context import ContextRepo
+    from faststream.nats import NatsBroker
+
+    broker = NatsBroker(context=ContextRepo({"data": "BROKER"}))
+    app = FastStream(broker, context=ContextRepo({"data": "APP"}))
+
+    @broker.subscriber("queue")
+    async def handle(data = Context()) -> None:
+        assert data == "BROKER"
+    ```
+
 ## Local
 
 To set a local context (available only within the message processing scope), use the context manager `scope`. It could me extremely uselful to fill context with additional options in [Middlewares](../middlewares/){.internal-link}

faststream/_internal/context/composition.py

@@ -0,0 +1,96 @@
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, TypeVar
+
+from typing_extensions import override
+
+from faststream._internal.constants import EMPTY
+from faststream._internal.context.repository import ContextRepo
+from faststream.exceptions import ContextError
+
+if TYPE_CHECKING:
+    from contextvars import Token
+
+T = TypeVar("T")
+
+
+class ContextRepoComposition(ContextRepo):
+    def __init__(self, *contexts: ContextRepo) -> None:
+        super().__init__()
+
+        self._inner_context = ContextRepo()
+        self._inner_context.set_global("context", self)
+
+        self._contexts = (self._inner_context, *contexts)
+
+    @property
+    @override
+    def context(self) -> dict[str, Any]:
+        result_context: dict[str, Any] = {}
+        for context in reversed(self._contexts):
+            result_context |= context.context
+
+        return result_context
+
+    @override
+    def set_global(self, key: str, v: Any) -> None:
+        self._inner_context.set_global(key, v)
+
+    @override
+    def reset_global(self, key: str) -> None:
+        self._inner_context.reset_global(key)
+
+    @override
+    def set_local(self, key: str, value: T) -> "Token[T]":
+        return self._inner_context.set_local(key, value)
+
+    @override
+    def reset_local(self, key: str, tag: "Token[Any]") -> None:
+        self._inner_context.reset_local(key, tag)
+
+    @override
+    def get_local(self, key: str, default: Any = None) -> Any:
+        for context in self._contexts:
+            variable = context.get_local(key, EMPTY)
+            if variable != EMPTY:
+                return variable
+        return default
+
+    @contextmanager
+    @override
+    def scope(self, key: str, value: Any) -> Generator[None, None, None]:
+        with self._inner_context.scope(key, value):
+            yield
+
+    @override
+    def get(self, key: str, default: Any = None) -> Any:
+        for context in self._contexts:
+            variable = context.get(key, EMPTY)
+            if variable != EMPTY:
+                return variable
+        return default
+
+    @override
+    def __getattr__(self, name: str, /) -> Any:
+        for context in self._contexts:
+            variable = getattr(context, name)
+            if variable is not None:
+                return variable
+        return None
+
+    @override
+    def resolve(self, argument: str) -> Any:
+        first, *_ = argument.split(".")
+
+        for context in self._contexts:
+            try:
+                return context.resolve(argument)
+            except ContextError:  # noqa: PERF203
+                pass
+
+        raise ContextError(self.context, first)
+
+    @override
+    def clear(self) -> None:
+        self._inner_context.clear()
+        self._inner_context.set_global("context", self)

faststream/_internal/context/repository.py

@@ -1,11 +1,13 @@
 from collections.abc import Generator, Mapping
 from contextlib import contextmanager
 from contextvars import ContextVar, Token
-from typing import Any
+from typing import Any, TypeVar
 
 from faststream._internal.constants import EMPTY
 from faststream.exceptions import ContextError
 
+T = TypeVar("T")
+
 
 class ContextRepo:
     """A class to represent a context repository."""
@@ -50,7 +52,7 @@ def reset_global(self, key: str) -> None:
         """
         self._global_context.pop(key, None)
 
-    def set_local(self, key: str, value: Any) -> "Token[Any]":
+    def set_local(self, key: str, value: T) -> "Token[T]":
         """Set a local context variable.
 
         Args:

faststream/_internal/di/config.py

@@ -8,6 +8,7 @@
 
 from faststream._internal.constants import EMPTY
 from faststream._internal.context import ContextRepo
+from faststream._internal.context.composition import ContextRepoComposition
 from faststream._internal.utils import apply_types, to_async
 
 if TYPE_CHECKING:
@@ -55,7 +56,7 @@ def __or__(self, value: "FastDependsConfig", /) -> "FastDependsConfig":
             use_fastdepends=use_fd,
             provider=value.provider,
             serializer=self.serializer or value.serializer,
-            context=self.context,
+            context=ContextRepoComposition(value.context, self.context),
             call_decorators=(*value.call_decorators, *self.call_decorators),
             get_dependent=self.get_dependent or value.get_dependent,
         )

faststream/_internal/fastapi/router.py

@@ -28,6 +28,7 @@
 from starlette.requests import Request
 from starlette.responses import JSONResponse, Response
 from starlette.routing import BaseRoute, _DefaultLifespan
+from typing_extensions import override
 
 from faststream._internal.application import StartAbleApplication
 from faststream._internal.broker import BrokerRouter
@@ -187,6 +188,11 @@ def __init__(
 
         self._lifespan_started = False
 
+    @property
+    @override
+    def context(self) -> ContextRepo:
+        return self.broker.context
+
     def _subscriber_compatibility_wrapper(
         self,
         dependencies: Iterable["params.Depends"] = (),

tests/brokers/base/consume.py

@@ -6,6 +6,7 @@
 from pydantic import BaseModel
 
 from faststream import Context, Depends, FastStream, TestApp
+from faststream.context import ContextRepo
 from faststream.exceptions import StopConsume
 
 from .basic import BaseTestcaseConfig
@@ -404,6 +405,123 @@ async def t() -> None: ...
             with pytest.raises(AssertionError):
                 await subscriber.get_one(timeout=1e-24)
 
+    async def test_composition_context(
+        self,
+        queue: str,
+        mock: MagicMock,
+        event: asyncio.Event,
+    ) -> None:
+        broker = self.get_broker(
+            apply_types=True,
+            context=ContextRepo({"broker_context": "broker_context"}),
+        )
+
+        app = FastStream(broker, context=ContextRepo({"app_context": "app_context"}))
+
+        args, kwargs = self.get_subscriber_params(queue)
+
+        @broker.subscriber(*args, **kwargs)
+        async def handle(app_context=Context(), broker_context=Context()) -> None:
+            mock(app_context, broker_context)
+            event.set()
+
+        async with self.patch_broker(broker) as br, TestApp(app):
+            await asyncio.wait(
+                (
+                    asyncio.create_task(br.publish("", queue)),
+                    asyncio.create_task(event.wait()),
+                ),
+                timeout=self.timeout,
+            )
+
+            assert event.is_set()
+            mock.assert_called_once_with("app_context", "broker_context")
+
+    async def test_composition_context_merge(
+        self,
+        queue: str,
+        mock: MagicMock,
+        event: asyncio.Event,
+    ) -> None:
+        broker = self.get_broker(
+            apply_types=True,
+            context=ContextRepo({"context_var": "BROKER"}),
+        )
+
+        app = FastStream(broker, context=ContextRepo({"context_var": "APP"}))
+
+        args, kwargs = self.get_subscriber_params(queue)
+
+        @broker.subscriber(*args, **kwargs)
+        async def handle(context_var=Context()) -> None:
+            mock(context_var)
+            event.set()
+
+        async with self.patch_broker(broker) as br, TestApp(app):
+            await asyncio.wait(
+                (
+                    asyncio.create_task(br.publish("", queue)),
+                    asyncio.create_task(event.wait()),
+                ),
+                timeout=self.timeout,
+            )
+
+            assert event.is_set()
+            mock.assert_called_once_with("BROKER")
+
+    async def test_multi_broker_composition_context(
+        self,
+        queue: str,
+        mock: MagicMock,
+        mock2: MagicMock,
+        event: asyncio.Event,
+        event2: asyncio.Event,
+    ) -> None:
+        broker = self.get_broker(
+            apply_types=True,
+            context=ContextRepo({"broker_var": "BROKER 1"}),
+        )
+        broker2 = self.get_broker(
+            apply_types=True,
+            context=ContextRepo({"broker_var": "BROKER 2"}),
+        )
+
+        app = FastStream(broker, broker2, context=ContextRepo({"app_var": "APP"}))
+
+        args, kwargs = self.get_subscriber_params(queue)
+
+        @broker.subscriber(*args, **kwargs)
+        async def handle(app_var=Context(), broker_var=Context()) -> None:
+            event.set()
+            mock(app_var, broker_var)
+
+        args2, kwargs2 = self.get_subscriber_params(queue + "2")
+
+        @broker2.subscriber(*args2, **kwargs2)
+        async def handle2(app_var=Context(), broker_var=Context()) -> None:
+            event2.set()
+            mock2(app_var, broker_var)
+
+        async with (
+            self.patch_broker(broker) as br,
+            self.patch_broker(broker2) as br2,
+            TestApp(app),
+        ):
+            await asyncio.wait(
+                (
+                    asyncio.create_task(br.publish("", queue)),
+                    asyncio.create_task(br2.publish("", queue + "2")),
+                    asyncio.create_task(event.wait()),
+                    asyncio.create_task(event2.wait()),
+                ),
+                timeout=self.timeout,
+            )
+
+        assert event.is_set()
+        assert event2.is_set()
+        mock.assert_called_once_with("APP", "BROKER 1")
+        mock2.assert_called_once_with("APP", "BROKER 2")
+
 
 @pytest.mark.asyncio()
 class BrokerRealConsumeTestcase(BrokerConsumeTestcase):