Update return types for task_with_context decorator

Author: CCXLV

python/fluxqueue/client.py

@@ -169,16 +169,16 @@ async def create_user_task(ctx: DbContext, email: str, username: str):
         """
 
         @overload
-        def decorator(func: Callable[Concatenate[C, P], None]) -> Callable[P, None]: ...
+        def decorator(func: Callable[Concatenate[C, P], None]) -> Task[P, None]: ...
 
         @overload
         def decorator(
             func: Callable[Concatenate[C, P], Coroutine[Any, Any, None]],
-        ) -> Callable[P, Coroutine[Any, Any, None]]: ...
+        ) -> Task[P, Coroutine[Any, Any, None]]: ...
 
         def decorator(
             func: Callable[Concatenate[C, P], None | Coroutine[Any, Any, None]],
-        ) -> Callable[P, None | Coroutine[Any, Any, None]]:
+        ) -> Task[P, None | Coroutine[Any, Any, None]]:
             return _with_context(
                 cast(Any, func),
                 name=name,

tests/test_schedule.py

@@ -1,3 +1,5 @@
+from fluxqueue import Context
+
 from .conftest import TestEnvFixture
 
 
@@ -8,3 +10,10 @@ def task():
 
     assert hasattr(task, "defer")
     assert hasattr(task, "cron")
+
+    @test_env.fluxqueue.task_with_context()
+    def task_with_ctx(ctx: Context):
+        pass
+
+    assert hasattr(task_with_ctx, "defer")
+    assert hasattr(task_with_ctx, "cron")