diff --git a/src/pydase/task/decorator.py b/src/pydase/task/decorator.py index a1d9a70..cce84c9 100644 --- a/src/pydase/task/decorator.py +++ b/src/pydase/task/decorator.py @@ -1,5 +1,3 @@ -import asyncio -import functools import logging from collections.abc import Callable, Coroutine from typing import Any, TypeVar @@ -13,18 +11,17 @@ R = TypeVar("R") def task( *, autostart: bool = False -) -> Callable[[Callable[[Any], Coroutine[None, None, R]]], Task[R]]: +) -> Callable[ + [ + Callable[[Any], Coroutine[None, None, R]] + | Callable[[], Coroutine[None, None, R]] + ], + Task[R], +]: def decorator( - func: Callable[[Any], Coroutine[None, None, R]], + func: Callable[[Any], Coroutine[None, None, R]] + | Callable[[], Coroutine[None, None, R]], ) -> Task[R]: - @functools.wraps(func) - async def wrapper(self: Any) -> R | None: - try: - return await func(self) - except asyncio.CancelledError: - logger.info("Task '%s' was cancelled", func.__name__) - return None - - return Task(wrapper, autostart=autostart) + return Task(func, autostart=autostart) return decorator diff --git a/src/pydase/task/task.py b/src/pydase/task/task.py index d471ccb..ba14473 100644 --- a/src/pydase/task/task.py +++ b/src/pydase/task/task.py @@ -96,7 +96,11 @@ class Task(pydase.data_service.data_service.DataService, Generic[R]): logger.info("Starting task %r", self._func_name) self._status = TaskStatus.RUNNING res: Coroutine[None, None, R] = self._bound_func() - return await res + try: + return await res + except asyncio.CancelledError: + logger.info("Task '%s' was cancelled", self._func_name) + return None logger.warning( "Cannot start task %r. Function has not been bound yet", self._func_name )