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 75% rename from python/fluxqueue/core.py rename to python/fluxqueue/client.py index 254cc2c..f211e94 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, 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) 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()