refactor(task): Improve task decorator type inference (#113)

* refactor(types): improve task decorator type inference

* fix(task): Raise TypeError for invalid return type annotations

Author: CCXLV

python/fluxqueue/__init__.py

@@ -1,3 +1,3 @@
 __all__ = ["FluxQueue"]
 
-from .core import FluxQueue
+from .client import FluxQueue

python/fluxqueue/core.py python/fluxqueue/client.pypython/fluxqueue/core.py renamed to python/fluxqueue/client.py

@@ -1,13 +1,12 @@
 import inspect
-from collections.abc import Callable
+from collections.abc import Callable, Coroutine
 from functools import wraps
-from typing import Any, ParamSpec, TypeAlias, cast
+from typing import Any, ParamSpec, cast, get_type_hints, overload
 
 from ._core import FluxQueueCore
 from .utils import get_task_name
 
 P = ParamSpec("P")
-TaskDecorator: TypeAlias = Callable[[Callable[P, Any]], Callable[P, Any]]
 
 
 class FluxQueue:
@@ -30,7 +29,7 @@ def task(
         name: str | None = None,
         queue: str = "default",
         max_retries: int = 3,
-    ) -> TaskDecorator[P]:
+    ):
         """
         Mark a function as a FluxQueue task.
 
@@ -52,7 +51,23 @@ def task(
             before treating it as dead.
         """
 
-        def decorator(func: Callable[P, Any]) -> Callable[P, Any]:
+        @overload
+        def decorator(func: Callable[P, None]) -> Callable[P, None]: ...
+
+        @overload
+        def decorator(
+            func: Callable[P, Coroutine[Any, Any, None]],
+        ) -> Callable[P, Coroutine[Any, Any, None]]: ...
+
+        def decorator(
+            func: Callable[P, None | Coroutine[Any, Any, None]],
+        ) -> Callable[P, None | Coroutine[Any, Any, None]]:
+            type_hints = get_type_hints(func)
+            return_type = type_hints.get("return")
+
+            if return_type and return_type is not type(None):
+                raise TypeError(f"Task function must return None, got {return_type}")
+
             is_async = inspect.iscoroutinefunction(func)
             task_name = get_task_name(func, name)
 

tests/test_tasks.py

@@ -30,3 +30,13 @@ async def async_hello(name: str):
     assert b"Async George" in redis_result[0]  # type: ignore
 
     test_env.redis_client.flushdb()
+
+
+def test_invalid_return_type(test_env: TestEnvFixture):
+    with pytest.raises(TypeError):
+
+        @test_env.fluxqueue.task()  # type: ignore
+        def test_task() -> int:
+            return 5
+
+        test_task()