From 2090ae3307e3bd9457af1ba0a22762219c2fbe09 Mon Sep 17 00:00:00 2001
From: Giorgi Merebashvili <mereba2627@gmail.com>
Date: Fri, 20 Feb 2026 15:57:27 +0400
Subject: [PATCH 1/2] refactor(types): improve task decorator type inference

---
 python/fluxqueue/__init__.py            |  2 +-
 python/fluxqueue/{core.py => client.py} | 19 ++++++++++++++-----
 2 files changed, 15 insertions(+), 6 deletions(-)
 rename python/fluxqueue/{core.py => client.py} (82%)

diff --git a/python/fluxqueue/__init__.py b/python/fluxqueue/__init__.py
index 2dc590a..60393a9 100644
--- a/python/fluxqueue/__init__.py
+++ b/python/fluxqueue/__init__.py
@@ -1,3 +1,3 @@
 __all__ = ["FluxQueue"]
 
-from .core import FluxQueue
+from .client import FluxQueue
diff --git a/python/fluxqueue/core.py b/python/fluxqueue/client.py
similarity index 82%
rename from python/fluxqueue/core.py
rename to python/fluxqueue/client.py
index 254cc2c..e5bda9d 100644
--- a/python/fluxqueue/core.py
+++ b/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, 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,17 @@ 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]]:
             is_async = inspect.iscoroutinefunction(func)
             task_name = get_task_name(func, name)
 

From df616a98181de4db38c85e32a3991721df20135a Mon Sep 17 00:00:00 2001
From: Giorgi Merebashvili <mereba2627@gmail.com>
Date: Fri, 20 Feb 2026 16:17:18 +0400
Subject: [PATCH 2/2] fix(task): Raise TypeError for invalid return type
 annotations

---
 python/fluxqueue/client.py |  8 +++++++-
 tests/test_tasks.py        | 10 ++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/python/fluxqueue/client.py b/python/fluxqueue/client.py
index e5bda9d..f211e94 100644
--- a/python/fluxqueue/client.py
+++ b/python/fluxqueue/client.py
@@ -1,7 +1,7 @@
 import inspect
 from collections.abc import Callable, Coroutine
 from functools import wraps
-from typing import Any, ParamSpec, cast, overload
+from typing import Any, ParamSpec, cast, get_type_hints, overload
 
 from ._core import FluxQueueCore
 from .utils import get_task_name
@@ -62,6 +62,12 @@ def decorator(
         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)
 
diff --git a/tests/test_tasks.py b/tests/test_tasks.py
index cbb425c..abfac0f 100644
--- a/tests/test_tasks.py
+++ b/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()