refactor: replace 'self: Task' magic with 'CURRENT_TASK' to fix issues with IDE typing

Author: malvex

src/sheppy/__init__.py

@@ -1,5 +1,5 @@
 from .backend import Backend, BackendError, MemoryBackend, RedisBackend
-from .models import Task
+from .models import CURRENT_TASK, Task
 from .queue import Queue
 from .task_factory import task
 from .testqueue import TestQueue
@@ -12,7 +12,7 @@
     # fastapi
     "Depends",
     # task
-    "task", "Task",
+    "task", "Task", "CURRENT_TASK",
     # queue
     "Queue",
     # testqueue

src/sheppy/models.py

@@ -26,6 +26,9 @@
 
 TASK_CRON_NS = UUID('7005b432-c135-4131-b19e-d3dc89703a9a')
 
+# sentinel object for current task injection (def my_task(x: int, task: Task = CURRENT_TASK): ...)
+CURRENT_TASK = object()
+
 
 def cron_expression_validator(value: str) -> str:
     if not croniter.is_valid(value):

src/sheppy/utils/task_execution.py

@@ -15,7 +15,7 @@
 import anyio
 from pydantic import ConfigDict, PydanticSchemaGenerationError, TypeAdapter
 
-from ..models import Task
+from ..models import CURRENT_TASK, Task
 from .fastapi import Depends
 
 cache_main_module: str | None = None
@@ -214,9 +214,13 @@ async def process_function_parameters(
         remaining_args = list(args)
 
         for param_name, param in list(signature.parameters.items()):
-            # Task injection (self: Task)
+            # current Task injection (current: Task = CURRENT_TASK)
             if task and TaskProcessor._is_task_injection(param):
-                final_args.append(task)
+                # inject positionally for positional params to maintain correct order
+                if param.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD):
+                    final_args.append(task)
+                else:
+                    final_kwargs[param_name] = task
                 continue
 
             # validate positional args
@@ -237,13 +241,7 @@ async def process_function_parameters(
 
     @staticmethod
     def _is_task_injection(param: inspect.Parameter) -> bool:
-        if param.name != 'self':
-            return False
-
-        if param.annotation == inspect.Parameter.empty:
-            return False
-
-        return param.annotation is Task or param.annotation == 'Task'
+        return param.default is CURRENT_TASK
 
 
     @staticmethod

src/sheppy/utils/validation.py

@@ -8,7 +8,7 @@
 
 from pydantic import TypeAdapter, ValidationError
 
-from sheppy.models import Task
+from sheppy.models import CURRENT_TASK
 
 from .fastapi import Depends
 
@@ -39,11 +39,8 @@ def validate_input(
 
     for param_name, param in signature.parameters.items():
 
+        # task self injection
         if _is_task_injection(param):
-            if param.default != inspect.Parameter.empty:
-                raise ValidationError.from_exception_data(
-                    f"Task injection parameter '{param_name}' cannot have a default value", line_errors=[]
-                )
             if param_name in remaining_kwargs:
                 raise ValidationError.from_exception_data(
                     f"Cannot provide value for Task injection parameter '{param_name}'", line_errors=[]
@@ -109,10 +106,7 @@ def validate_input(
 
 
 def _is_task_injection(param: inspect.Parameter) -> bool:
-    if param.name != 'self':
-        return False
-
-    return param.annotation is Task or param.annotation == 'Task'
+    return param.default is CURRENT_TASK
 
 
 def _is_depends_parameter(param: inspect.Parameter) -> bool:

tests/conftest.py

@@ -4,7 +4,7 @@
 import pytest
 import pytest_asyncio
 
-from sheppy import Queue, Task, Worker, task
+from sheppy import CURRENT_TASK, Queue, Task, Worker, task
 from sheppy.backend import Backend, MemoryBackend, RedisBackend
 
 TEST_QUEUE_NAME = "pytest"
@@ -67,15 +67,15 @@ async def queue(backend: Backend) -> Queue:
 
 
 @task(retry=2, retry_delay=0.1)
-async def async_fail_once(self: Task) -> str:
-    if self.retry_count == 0:
+async def async_fail_once(current: Task = CURRENT_TASK) -> str:
+    if current.retry_count == 0:
         raise Exception("transient error")
     return "ok"
 
 
 @task(retry=2, retry_delay=0)
-def sync_fail_once(self: Task) -> str:
-    if self.retry_count == 0:
+def sync_fail_once(current: Task = CURRENT_TASK) -> str:
+    if current.retry_count == 0:
         raise Exception("transient error")
     return "ok"
 

tests/dependencies.py

@@ -5,7 +5,7 @@
 
 from pydantic import BaseModel
 
-from sheppy import Depends, Task, task
+from sheppy import CURRENT_TASK, Depends, Task, task
 
 
 class User(BaseModel):
@@ -139,18 +139,18 @@ async def async_task_with_pydantic_model(user: User) -> Status:
 
 
 @task
-def task_with_self(self: Task, x: int, y: int) -> dict[str, Any]:
-    return {"task_id": self.id, "result": x + y}
+def task_with_current_task(current: Task = CURRENT_TASK, x: int = 5, y: int = 6) -> dict[str, Any]:
+    return {"task_id": current.id, "result": x + y}
 
 
 @task
-def task_with_self_middle(x: int, self: Task, y: int) -> dict[str, Any]:
-    return {"task_id": self.id, "result": x + y}
+def task_with_current_task_middle(x: int, current: Task = CURRENT_TASK, y: int = 7) -> dict[str, Any]:
+    return {"task_id": current.id, "result": x + y}
 
 
 @task
-def task_with_self_end(x: int, y: int, self: Task) -> dict[str, Any]:
-    return {"task_id": self.id, "result": x + y}
+def task_with_current_task_end(x: int, y: int, current: Task = CURRENT_TASK) -> dict[str, Any]:
+    return {"task_id": current.id, "result": x + y}
 
 
 @task
@@ -482,9 +482,9 @@ def deep_recursion_tasks() -> list[TaskTestCase]:
     def self_referencing_tasks() -> list[TaskTestCase]:
         """Tasks that should fail."""
         return [
-            TaskTestCase("task_with_self", task_with_self, (22, 33), expected_result=55),
-            TaskTestCase("task_with_self_middle", task_with_self_middle, (22, 33), expected_result=55),
-            TaskTestCase("task_with_self_end", task_with_self_end, (22, 33), expected_result=55),
+            TaskTestCase("task_with_current_task", task_with_current_task, (), {"x": 22, "y": 33}, expected_result=55),
+            TaskTestCase("task_with_current_task_middle", task_with_current_task_middle, (22, ), {"y": 33}, expected_result=55),
+            TaskTestCase("task_with_current_task_end", task_with_current_task_end, (22, 33), expected_result=55),
         ]
 
     @staticmethod