Added context injector. (#38)

Signed-off-by: Pavel Kirilin <[email protected]>

Author: s3rius

taskiq/brokers/inmemory_broker.py

@@ -124,13 +124,10 @@ async def kick(self, message: BrokerMessage) -> None:
         target_task = self.available_tasks.get(message.task_name)
         if target_task is None:
             raise TaskiqError("Unknown task.")
-        if self.receiver.task_signatures:
-            if not self.receiver.task_signatures.get(target_task.task_name):
-                self.receiver.task_signatures[
-                    target_task.task_name
-                ] = inspect.signature(
-                    target_task.original_func,
-                )
+        if not self.receiver.task_signatures.get(target_task.task_name):
+            self.receiver.task_signatures[target_task.task_name] = inspect.signature(
+                target_task.original_func,
+            )
 
         await self.receiver.callback(message=message)
 

taskiq/cli/receiver.py

@@ -4,21 +4,48 @@
 from concurrent.futures import ThreadPoolExecutor
 from logging import getLogger
 from time import time
-from typing import Any, Callable, Dict
+from typing import Any, Callable, Dict, Optional
 
 from taskiq.abc.broker import AsyncBroker
 from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.cli.args import TaskiqArgs
 from taskiq.cli.log_collector import log_collector
 from taskiq.cli.params_parser import parse_params
-from taskiq.context import Context, context_updater
+from taskiq.context import Context
 from taskiq.message import BrokerMessage, TaskiqMessage
 from taskiq.result import TaskiqResult
 from taskiq.utils import maybe_awaitable
 
 logger = getLogger(__name__)
 
 
+def inject_context(
+    signature: Optional[inspect.Signature],
+    message: TaskiqMessage,
+    broker: AsyncBroker,
+) -> None:
+    """
+    Inject context parameter in message's kwargs.
+
+    This function parses signature to get
+    the context parameter definition.
+
+    If at least one parameter has the Context
+    type, it will add current context as kwarg.
+
+    :param signature: function's signature.
+    :param message: current taskiq message.
+    :param broker: current broker.
+    """
+    if signature is None:
+        return
+    for param_name, param in signature.parameters.items():
+        if param.annotation is param.empty:
+            continue
+        if param.annotation is Context:
+            message.kwargs[param_name] = Context(message.copy(), broker)
+
+
 def _run_sync(target: Callable[..., Any], message: TaskiqMessage) -> Any:
     """
     Runs function synchronously.
@@ -40,11 +67,8 @@ def __init__(self, broker: AsyncBroker, cli_args: TaskiqArgs) -> None:
         self.broker = broker
         self.cli_args = cli_args
         self.task_signatures: Dict[str, inspect.Signature] = {}
-        if not cli_args.no_parse:
-            for task in self.broker.available_tasks.values():
-                self.task_signatures[task.task_name] = inspect.signature(
-                    task.original_func,
-                )
+        for task in self.broker.available_tasks.values():
+            self.task_signatures[task.task_name] = inspect.signature(task.original_func)
         self.executor = ThreadPoolExecutor(
             max_workers=cli_args.max_threadpool_threads,
         )
@@ -100,11 +124,10 @@ async def callback(  # noqa: C901
             taskiq_msg.task_name,
             taskiq_msg.task_id,
         )
-        with context_updater(Context(taskiq_msg, self.broker)):
-            result = await self.run_task(
-                target=self.broker.available_tasks[message.task_name].original_func,
-                message=taskiq_msg,
-            )
+        result = await self.run_task(
+            target=self.broker.available_tasks[message.task_name].original_func,
+            message=taskiq_msg,
+        )
         for middleware in self.broker.middlewares:
             if middleware.__class__.post_execute != TaskiqMiddleware.post_execute:
                 await maybe_awaitable(middleware.post_execute(taskiq_msg, result))
@@ -147,7 +170,15 @@ async def run_task(  # noqa: C901, WPS210
         logs = io.StringIO()
         returned = None
         found_exception = None
-        parse_params(self.task_signatures.get(message.task_name), message)
+        signature = self.task_signatures.get(message.task_name)
+        if self.cli_args.no_parse:
+            signature = None
+        parse_params(signature, message)
+        inject_context(
+            self.task_signatures.get(message.task_name),
+            message,
+            self.broker,
+        )
         # Captures function's logs.
         with log_collector(logs, self.cli_args.log_collector_format):
             # Start a timer.

taskiq/cli/tests/test_context.py

@@ -0,0 +1,79 @@
+import inspect
+
+from taskiq.cli.receiver import inject_context
+from taskiq.context import Context
+from taskiq.message import TaskiqMessage
+
+
+def test_inject_context_success() -> None:
+    """Test that context variable is injected as expected."""
+
+    def func(param1: int, ctx: Context) -> int:
+        return param1
+
+    message = TaskiqMessage(
+        task_id="",
+        task_name="",
+        labels={},
+        args=[1],
+        kwargs={},
+    )
+
+    inject_context(
+        inspect.signature(func),
+        message=message,
+        broker=None,  # type: ignore
+    )
+
+    assert message.kwargs.get("ctx")
+    assert isinstance(message.kwargs["ctx"], Context)
+
+
+def test_inject_context_no_typehint() -> None:
+    """Test that context won't be injected in untyped parameter."""
+
+    def func(param1: int, ctx) -> int:  # type: ignore
+        return param1
+
+    message = TaskiqMessage(
+        task_id="",
+        task_name="",
+        labels={},
+        args=[1],
+        kwargs={},
+    )
+
+    inject_context(
+        inspect.signature(func),
+        message=message,
+        broker=None,  # type: ignore
+    )
+
+    assert message.kwargs.get("ctx") is None
+
+
+def test_inject_context_no_ctx_parameter() -> None:
+    """
+    Tests that injector won't raise an error.
+
+    If the Context-typed parameter doesn't exist.
+    """
+
+    def func(param1: int) -> int:
+        return param1
+
+    message = TaskiqMessage(
+        task_id="",
+        task_name="",
+        labels={},
+        args=[1],
+        kwargs={},
+    )
+
+    inject_context(
+        inspect.signature(func),
+        message=message,
+        broker=None,  # type: ignore
+    )
+
+    assert not message.kwargs

taskiq/context.py

@@ -1,6 +1,3 @@
-from contextlib import contextmanager
-from typing import Generator
-
 from taskiq.abc.broker import AsyncBroker
 from taskiq.message import TaskiqMessage
 
@@ -14,41 +11,3 @@ def __init__(self, message: TaskiqMessage, broker: AsyncBroker) -> None:
 
 
 default_context = Context(None, None)  # type: ignore
-current_context = None
-
-
-@contextmanager
-def context_updater(new_context: Context) -> Generator[None, None, None]:
-    """
-    Update context for some time.
-
-    :param new_context: new context to set.
-    :yield: nothing.
-    """
-    global current_context  # noqa: WPS420
-    current_context = new_context  # noqa: WPS442
-
-    yield
-
-    current_context = None  # noqa: WPS442
-
-
-def get_context() -> Context:
-    """
-    Get current context.
-
-    This function always return contexts,
-    but if you call this function inside tests,
-    or somewhere you have to be careful,
-    since if current_context is None it will
-    return default_context.
-
-    To override context please use context_updater
-    context manager.
-
-    :return: context.
-    """
-    global current_context  # noqa: WPS420
-    if current_context is None:
-        return default_context
-    return current_context