diff --git a/aioclock/app.py b/aioclock/app.py index 0422c23..3ddd218 100644 --- a/aioclock/app.py +++ b/aioclock/app.py @@ -1,6 +1,12 @@ import asyncio +import sys from functools import wraps -from typing import Any, Callable +from typing import Any, Awaitable, Callable, TypeVar + +if sys.version_info < (3, 10): + from typing_extensions import ParamSpec +else: + from typing import ParamSpec from fast_depends import inject @@ -10,6 +16,9 @@ from aioclock.types import Triggers from aioclock.utils import flatten_chain +T = TypeVar("T") +P = ParamSpec("P") + class AioClock: """ @@ -117,9 +126,9 @@ async def main(): ``` """ - def decorator(func: Callable[..., Any]) -> Callable[..., Any]: + def decorator(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]: @wraps(func) - async def wrapper(*args, **kwargs) -> Any: + async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T: return await func(*args, **kwargs) self._app_tasks.append( diff --git a/aioclock/group.py b/aioclock/group.py index f94a0ea..a59f74d 100644 --- a/aioclock/group.py +++ b/aioclock/group.py @@ -1,6 +1,12 @@ import asyncio +import sys from functools import wraps -from typing import Any, Callable, Union +from typing import Awaitable, Callable, TypeVar, Union + +if sys.version_info < (3, 10): + from typing_extensions import ParamSpec +else: + from typing import ParamSpec from fast_depends import inject @@ -8,6 +14,9 @@ from aioclock.task import Task from aioclock.triggers import BaseTrigger +T = TypeVar("T") +P = ParamSpec("P") + class Group: def __init__(self, *, tasks: Union[list[Task], None] = None): @@ -50,9 +59,9 @@ async def send_email(): ``` """ - def decorator(func: Callable[..., Any]) -> Callable[..., Any]: + def decorator(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]: @wraps(func) - async def wrapper(*args, **kwargs) -> Any: + async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T: return await func(*args, **kwargs) self._tasks.append( diff --git a/aioclock/task.py b/aioclock/task.py index b437caf..284f01e 100644 --- a/aioclock/task.py +++ b/aioclock/task.py @@ -1,5 +1,5 @@ from dataclasses import dataclass -from typing import Any, Callable +from typing import Any, Awaitable, Callable from aioclock.logger import logger from aioclock.triggers import BaseTrigger @@ -12,7 +12,7 @@ class Task: This is internally used, when you decorate your function with `aioclock.task`. """ - func: Callable[..., Any] + func: Callable[..., Awaitable[Any]] """Decorated function that will be run by AioClock.""" trigger: BaseTrigger