diff --git a/src/resonate/resonate.py b/src/resonate/resonate.py index 911c9e1c..f9441c65 100644 --- a/src/resonate/resonate.py +++ b/src/resonate/resonate.py @@ -181,7 +181,7 @@ def register( (func, Options(version=version, durable=True, retry_policy=retry_policy)), ) - def fn( + def __call__( self, name: str | None = None, version: int = 1, @@ -189,7 +189,20 @@ def fn( ) -> Callable[ [Callable[Concatenate[Context, P], Any]], RegisteredFn[P, Any], - ]: ... + ]: + def wrapper( + func: Callable[Concatenate[Context, P], Any], + ) -> RegisteredFn[P, Any]: + self._registry.add( + name or func.__name__, + ( + func, + Options(version=version, durable=True, retry_policy=retry_policy), + ), + ) + return RegisteredFn(self._scheduler, func) + + return wrapper @overload def run( diff --git a/tests/test_functionality.py b/tests/test_functionality.py index a0c5cc75..bb7ec474 100644 --- a/tests/test_functionality.py +++ b/tests/test_functionality.py @@ -459,7 +459,7 @@ def test_golden_device_rfi_and_lfc_with_decorator() -> None: task_source=Poller("http://localhost:8002", group=group), ) - @resonate.fn() + @resonate() def foo(ctx: Context, n: str) -> Generator[Yieldable, Any, str]: v: str = yield ctx.lfc(bar, n).options( id="bar", @@ -472,7 +472,7 @@ def bar(ctx: Context, n: str) -> Generator[Yieldable, Any, str]: v: str = yield p return v - @resonate.fn() + @resonate() def baz(ctx: Context, n: str) -> str: # noqa: ARG001 return n @@ -611,7 +611,7 @@ def test_golden_device_rfc_and_lfc_with_decorator() -> None: task_source=Poller("http://localhost:8002", group=group), ) - @resonate.fn() + @resonate() def foo(ctx: Context, n: str) -> Generator[Yieldable, Any, str]: v: str = yield ctx.lfc(bar, n).options( id="bar", @@ -623,7 +623,7 @@ def bar(ctx: Context, n: str) -> Generator[Yieldable, Any, str]: v: str = yield ctx.rfc(baz, n).options(id="baz", send_to=poll(group)) return v - @resonate.fn() + @resonate() def baz(ctx: Context, n: str) -> str: # noqa: ARG001 return n @@ -734,7 +734,7 @@ def test_sleep() -> None: store=store, task_source=Poller("http://localhost:8002", group=group) ) - @resonate.fn() + @resonate() def foo_sleep(ctx: Context, n: int) -> Generator[Yieldable, Any, int]: yield ctx.sleep(n) return n @@ -794,7 +794,7 @@ def test_golden_device_detached_with_registered() -> None: task_source=Poller("http://localhost:8002", group=group), ) - @resonate.fn() + @resonate() def foo_golden_device_detached_with_registered( ctx: Context, n: str ) -> Generator[Yieldable, Any, str]: @@ -807,11 +807,11 @@ def foo_golden_device_detached_with_registered( v: str = yield p return v - @resonate.fn() + @resonate() def bar_golden_device_detached_with_registered(ctx: Context, n: str) -> str: # noqa: ARG001 return n - @resonate.fn() + @resonate() def baz_golden_device_detached_with_registered( ctx: Context, # noqa: ARG001 promise_id: str,