feat: improve decorator

Author: brian14708

pyproject.toml

@@ -48,7 +48,7 @@ future-annotations = true
 
 [tool.ruff.lint.flake8-type-checking]
 runtime-evaluated-base-classes = ["typing_extensions.TypedDict"]
-runtime-evaluated-decorators = ["duron.fn", "duron.host"]
+runtime-evaluated-decorators = ["duron.fn", "duron.effect"]
 
 [tool.ruff.lint.flake8-tidy-imports.banned-api]
 "typing.TypedDict".msg = "Use typing_extensions.TypedDict instead."

src/duron/__init__.py

@@ -1,7 +1,7 @@
 from .config import set_config as set_config
 from .context import Context as Context
+from .fn import effect as effect
 from .fn import fn as fn
-from .fn import host as host
 from .stream import EndOfStream as EndOfStream
 from .stream import Sink as Sink
 from .stream import Stream as Stream

src/duron/fn.py

@@ -8,6 +8,7 @@
     Generic,
     ParamSpec,
     TypeVar,
+    overload,
 )
 
 from duron.config import config
@@ -39,12 +40,28 @@ def create_task(self, log: LogStorage) -> TaskGuard[_P, _T_co]:
         return TaskGuard(Task(self, log))
 
 
+@overload
+def fn(
+    _fn: Callable[Concatenate[Context, _P], Coroutine[Any, Any, _T_co]],
+) -> Fn[_P, _T_co]: ...
+@overload
 def fn(
     *, codec: Codec | None = None
 ) -> Callable[
     [Callable[Concatenate[Context, _P], Coroutine[Any, Any, _T_co]]],
     Fn[_P, _T_co],
-]:
+]: ...
+def fn(
+    _fn: Callable[Concatenate[Context, _P], Coroutine[Any, Any, _T_co]] | None = None,
+    *,
+    codec: Codec | None = None,
+) -> (
+    Fn[_P, _T_co]
+    | Callable[
+        [Callable[Concatenate[Context, _P], Coroutine[Any, Any, _T_co]]],
+        Fn[_P, _T_co],
+    ]
+):
     """
     Mark a function as durable, meaning its execution can be recorded and
     replayed.
@@ -55,8 +72,23 @@ def decorate(
     ) -> Fn[_P, _T_co]:
         return Fn(codec=codec or config.codec, fn=fn)
 
-    return decorate
-
-
-def host(fn: Callable[_P, _T_co]) -> Callable[_P, _T_co]:
-    return fn
+    if _fn is not None:
+        return decorate(_fn)
+    else:
+        return decorate
+
+
+@overload
+def effect(_fn: Callable[_P, _T_co]) -> Callable[_P, _T_co]: ...
+@overload
+def effect() -> Callable[[Callable[_P, _T_co]], Callable[_P, _T_co]]: ...
+def effect(
+    _fn: Callable[_P, _T_co] | None = None,
+) -> Callable[_P, _T_co] | Callable[[Callable[_P, _T_co]], Callable[_P, _T_co]]:
+    def decorate(fn: Callable[_P, _T_co]) -> Callable[_P, _T_co]:
+        return fn
+
+    if _fn is not None:
+        return decorate(_fn)
+    else:
+        return decorate

tests/test_stream.py

@@ -6,7 +6,7 @@
 
 import pytest
 
-from duron import EndOfStream, Sink, fn, host
+from duron import EndOfStream, Sink, effect, fn
 from duron.context import Context
 from duron.contrib.storage import MemoryLogStorage
 
@@ -48,7 +48,7 @@ async def test_stream_host():
     async def activity(ctx: Context) -> None:
         stream, handle = await ctx.create_stream(int)
 
-        @host
+        @effect
         async def task(stream: Sink[int]):
             for i in range(50):
                 await stream.send(i)
@@ -70,7 +70,7 @@ async def test_run():
 
     @fn()
     async def activity(ctx: Context) -> None:
-        @host
+        @effect
         async def f(s: str) -> AsyncGenerator[str, str]:
             while len(s) < 100:
                 if len(s) == sleep_idx: