feat: use type annotations for stateful

Author: brian14708

examples/hello_world.py

@@ -30,11 +30,11 @@ async def generate_lucky_number() -> int:
     return random.randint(1, 100)
 
 
-@duron.effect(stateful=True, initial=lambda: 0, reducer=int.__add__)
+@duron.effect
 async def count_up(count: int, target: int) -> AsyncGenerator[int, int]:
     await asyncio.sleep(0.5)
     while count < target:
-        count = yield 10
+        count = yield (count + 10)
         logger.info("⚡ Current count: %s", count)
         await asyncio.sleep(0.05)
 
@@ -44,7 +44,7 @@ async def greeting_flow(ctx: duron.Context, name: str) -> str:
     message, lucky_number = await asyncio.gather(
         ctx.run(work, name), ctx.run(generate_lucky_number)
     )
-    _ = await ctx.run(count_up, lucky_number)
+    _ = await ctx.run(count_up, 0, lucky_number)
     return f"{message} Your lucky number is {lucky_number}."
 
 

src/duron/__init__.py

@@ -8,5 +8,6 @@
 from duron._core.stream import StreamClosed as StreamClosed
 from duron._core.stream import StreamWriter as StreamWriter
 from duron._decorator.durable import durable as durable
+from duron._decorator.effect import Reducer as Reducer
 from duron._decorator.effect import effect as effect
 from duron.typing import Provided as Provided

src/duron/_core/context.py

@@ -5,8 +5,8 @@
 import functools
 import inspect
 from random import Random
-from typing import TYPE_CHECKING, cast
-from typing_extensions import Any, ParamSpec, TypeVar, final, overload
+from typing import TYPE_CHECKING, Concatenate, cast, get_args, get_origin
+from typing_extensions import Any, AsyncGenerator, ParamSpec, TypeVar, final, overload
 
 from duron._core.ops import (
     Barrier,
@@ -18,8 +18,9 @@
 )
 from duron._core.signal import create_signal
 from duron._core.stream import create_stream, run_stateful
-from duron._decorator.effect import EffectFn, StatefulFn
+from duron._decorator.effect import Reducer
 from duron.typing import inspect_function
+from duron.typing._hint import UnspecifiedType
 
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable, Coroutine
@@ -46,28 +47,25 @@ def __init__(self, loop: EventLoop, seed: str) -> None:
     @overload
     async def run(
         self,
-        fn: Callable[_P, Coroutine[Any, Any, _T]] | EffectFn[_P, _T],
+        fn: Callable[_P, Coroutine[Any, Any, _T]],
         /,
         *args: _P.args,
         **kwargs: _P.kwargs,
     ) -> _T: ...
     @overload
     async def run(
         self,
-        fn: Callable[_P, _T] | StatefulFn[_P, _T, Any],
+        fn: Callable[Concatenate[_T, _P], AsyncGenerator[_S, _T]],
         /,
+        state: _S,
         *args: _P.args,
         **kwargs: _P.kwargs,
     ) -> _T: ...
+    @overload
     async def run(
-        self,
-        fn: Callable[_P, Coroutine[Any, Any, _T] | _T]
-        | EffectFn[_P, _T]
-        | StatefulFn[_P, _T, Any],
-        /,
-        *args: _P.args,
-        **kwargs: _P.kwargs,
-    ) -> _T:
+        self, fn: Callable[_P, _T], /, *args: _P.args, **kwargs: _P.kwargs
+    ) -> _T: ...
+    async def run(self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> Any:
         """
         Run a function within the context.
 
@@ -81,18 +79,13 @@ async def run(
             msg = "Context time can only be used in the context loop"
             raise RuntimeError(msg)
 
-        if isinstance(fn, StatefulFn):
-            async with self.stream(
-                cast("StatefulFn[_P, _T, Any]", fn), *args, **kwargs
-            ) as (stream, result):
+        if inspect.isasyncgenfunction(fn):
+            async with self.stream(fn, *args, **kwargs) as (stream, result):
                 await stream.discard()
                 return await result
 
         callable_: Callable[[], Coroutine[Any, Any, object]]
-        if isinstance(fn, EffectFn):
-            hint = fn.type_hint
-            callable_ = functools.partial(fn.fn, *args, **kwargs)
-        elif inspect.iscoroutinefunction(fn):
+        if inspect.iscoroutinefunction(fn):
             hint = inspect_function(fn)
             callable_ = functools.partial(fn, *args, **kwargs)
         else:
@@ -103,7 +96,7 @@ async def wrapper() -> object:  # noqa: RUF029
             hint = inspect_function(fn)
             callable_ = wrapper
 
-        op: asyncio.Future[_T] = create_op(
+        op: asyncio.Future[object] = create_op(
             self._loop,
             FnCall(
                 callable=callable_,
@@ -115,7 +108,12 @@ async def wrapper() -> object:  # noqa: RUF029
         return await op
 
     def stream(
-        self, fn: StatefulFn[_P, _T, _S], /, *args: _P.args, **kwargs: _P.kwargs
+        self,
+        fn: Callable[Concatenate[_T, _P], AsyncGenerator[_S, _T]],
+        /,
+        initial: _T,
+        *args: _P.args,
+        **kwargs: _P.kwargs,
     ) -> AbstractAsyncContextManager[tuple[Stream[_S], Awaitable[_T]]]:
         """Stream stateful function partial results.
 
@@ -133,8 +131,17 @@ def stream(
         if asyncio.get_running_loop() is not self._loop:
             msg = "Context time can only be used in the context loop"
             raise RuntimeError(msg)
+
+        type_hint = inspect_function(fn)
+        action_type: TypeHint[_S] = UnspecifiedType
+        if get_origin(ret := type_hint.return_type) is AsyncGenerator:
+            action_type, _ = get_args(ret)
+
+        state_name = type_hint.parameters[0]
+        annotations = type_hint.parameter_annotations.get(state_name, ())
+        reducer = _find_reducer(tuple(annotations))
         return run_stateful(
-            self._loop, fn.action_type, fn.initial(), fn.reducer, fn.fn, *args, **kwargs
+            self._loop, action_type, reducer, fn, initial, *args, **kwargs
         )
 
     async def create_stream(
@@ -288,3 +295,20 @@ async def complete_future(
             self._loop,
             FutureComplete(future_id=future_id, value=result, exception=exception),
         )
+
+
+def _find_reducer(annotations: tuple[Any, ...]) -> Callable[[_S, _T], _S]:
+    for annotation in annotations:
+        if not isinstance(annotation, Reducer):
+            continue
+        hint = inspect_function(annotation.reducer)
+        if len(hint.parameters) != 2:
+            msg = "Reducer function must have exactly two parameters"
+            raise TypeError(msg)
+        return cast("Callable[[_S, _T], _S]", annotation.reducer)
+
+    return cast("Callable[[_S, _T], _S]", _default_reducer)
+
+
+def _default_reducer(_old: object, new: object) -> object:
+    return new

src/duron/_core/session.py

@@ -40,6 +40,7 @@
 from duron.tracing._tracer import current_tracer, span
 from duron.typing import JSONValue
 from duron.typing._hint import UnspecifiedType
+from duron.typing._inspect import inspect_function
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Coroutine
@@ -770,7 +771,7 @@ async def _prelude_fn(
     ctx = Context(loop, init_params["nonce"])
 
     codec = fn.codec
-    type_info = fn.type_hints
+    type_info = inspect_function(fn.fn)
     args = tuple(
         codec.decode_json(
             arg,

src/duron/_core/stream.py

@@ -354,18 +354,18 @@ async def __aexit__(
 async def run_stateful(
     loop: EventLoop,
     dtype: TypeHint[Any],
-    initial: _T,
     reducer: Callable[[_T, _U], _T],
     fn: Callable[Concatenate[_T, _P], AsyncGenerator[_U, _T]],
     /,
+    initial: _T,
     *args: _P.args,
     **kwargs: _P.kwargs,
 ) -> AsyncGenerator[tuple[Stream[_U], Awaitable[_T]], None]:
     assert asyncio.get_running_loop() is loop
 
     name = cast("str", getattr(fn, "__name__", repr(fn)))
     stream: _StatefulStream[_U, _T] = _StatefulStream(
-        initial, reducer, fn, *args, **kwargs
+        reducer, fn, initial, *args, **kwargs
     )
     sink: StreamWriter[_U] = OpWriter(
         await create_op(
@@ -410,10 +410,10 @@ class _StatefulStream(_BufferStream[_U], Generic[_U, _T]):
 
     def __init__(
         self,
-        initial: _T,
         reducer: Callable[[_T, _U], _T],
         fn: Callable[Concatenate[_T, _P], AsyncGenerator[_U, _T]],
         /,
+        initial: _T,
         *args: _P.args,
         **kwargs: _P.kwargs,
     ) -> None:

src/duron/_decorator/durable.py

@@ -40,7 +40,6 @@ def __init__(
         self.codec = codec
         self.fn = fn
         self.inject = sorted(inject)
-        self.type_hints = inspect_function(fn)
 
 
 @overload

src/duron/_decorator/effect.py

@@ -1,137 +1,58 @@
 from __future__ import annotations
 
-import functools
-from collections.abc import AsyncGenerator
-from typing import TYPE_CHECKING, Concatenate, Generic, Literal, get_args, get_origin
-from typing_extensions import Any, ParamSpec, TypeVar, final, overload
-
-from duron.typing import UnspecifiedType, inspect_function
+from typing import TYPE_CHECKING
+from typing_extensions import NamedTuple, ParamSpec, TypeVar, overload
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Coroutine
-
-    from duron.typing import TypeHint
+    from collections.abc import Callable
 
 
-_T = TypeVar("_T")
-_S = TypeVar("_S")
 _T_co = TypeVar("_T_co", covariant=True)
 _P = ParamSpec("_P")
 
 
-@final
-class EffectFn(Generic[_P, _T_co]):
-    def __init__(self, fn: Callable[_P, Coroutine[Any, Any, _T_co]]) -> None:
-        self.fn = fn
-        self.type_hint = inspect_function(fn)
-        functools.update_wrapper(self, fn)
-
-    def __call__(
-        self, *args: _P.args, **kwargs: _P.kwargs
-    ) -> Coroutine[Any, Any, _T_co]:
-        return self.fn(*args, **kwargs)
-
-
-@final
-class StatefulFn(Generic[_P, _S, _T]):
-    def __init__(
-        self,
-        fn: Callable[Concatenate[_S, _P], AsyncGenerator[_T, _S]],
-        initial: Callable[[], _S],
-        reducer: Callable[[_S, _T], _S],
-    ) -> None:
-        self.fn = fn
-        self.type_hint = inspect_function(fn)
-        self.initial = initial
-        self.reducer = reducer
-
-        action_type: TypeHint[_T] = UnspecifiedType
-        if get_origin(ret := self.type_hint.return_type) is AsyncGenerator:
-            action_type, _ = get_args(ret)
-        self.action_type = action_type
-        functools.update_wrapper(self, fn)
-
-    def __call__(
-        self, state: _S, *args: _P.args, **kwargs: _P.kwargs
-    ) -> AsyncGenerator[_T, _S]:
-        return self.fn(state, *args, **kwargs)
+class Reducer(NamedTuple):
+    """Annotation to mark a parameter as a reducer."""
+
+    reducer: Callable[[object, object], object]
 
 
 @overload
-def effect(fn: Callable[_P, Coroutine[Any, Any, _T_co]], /) -> EffectFn[_P, _T_co]: ...
-@overload
-def effect(
-    *, stateful: Literal[False] = ...
-) -> Callable[[Callable[_P, Coroutine[Any, Any, _T_co]]], EffectFn[_P, _T_co]]: ...
+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(
-    *,
-    stateful: Literal[True],
-    initial: Callable[[], _S],
-    reducer: Callable[[_S, _T], _S],
-) -> Callable[
-    [Callable[Concatenate[_S, _P], AsyncGenerator[_T, _S]]], StatefulFn[_P, _S, _T]
-]: ...
-def effect(
-    fn: Callable[_P, Coroutine[Any, Any, _T_co]] | None = None,
-    /,
-    *,
-    # stateful parameters
-    stateful: bool = False,
-    initial: Callable[[], _S] | None = None,
-    reducer: Callable[[_S, _T], _S] | None = None,
-) -> (
-    EffectFn[_P, _T_co]
-    | Callable[[Callable[_P, Coroutine[Any, Any, _T_co]]], EffectFn[_P, _T_co]]
-    | Callable[
-        [Callable[Concatenate[_S, _P], AsyncGenerator[_T, _S]]], StatefulFn[_P, _S, _T]
-    ]
-):
+    fn: Callable[_P, _T_co] | None = None, /
+) -> Callable[_P, _T_co] | Callable[[Callable[_P, _T_co]], Callable[_P, _T_co]]:
     """Decorator to mark async functions as effects.
 
     Effects are operations that interact with the outside world.
 
-    Args:
-        stateful: Whether the effect is stateful.
-        initial: Factory function for initial state (required with `stateful=True`)
-        reducer: Function to reduce actions into state (required with `stateful=True`)
-
     Example:
-        Basic example:
         ```python
         @duron.effect
-        async def fetch_data(url: str) -> dict:
-            return await http_client.get(url)
-        ```
+        async def send_email(to: str, subject: str, body: str) -> None:
+            # Send an email
+            ...
 
-        Stateful example:
-        ```python
-        @duron.effect(stateful=True, initial=lambda: 0, reducer=int.__add__)
-        async def count_items(state: int, items: list) -> AsyncGenerator[int, int]:
-            # restore based on `state`
-            for item in items:
-                yield 1
+
+        @duron.effect
+        async def counter(
+            state: Annotated[int, duron.Reducer(lambda s, a: s + a)], increment: int
+        ) -> AsyncGenerator[int, int]:
+            state += increment
+            yield state
         ```
 
+
     Returns:
         Function wrapper that can be invoked with [ctx.run()][duron.Context.run]
-
-    Raises:
-        ValueError: If stateful is True but initial or reducer is not provided.
     """
-    if fn is not None:
-        return EffectFn(fn)
 
-    if stateful:
-        if not initial or not reducer:
-            msg = "initial and reducer must be provided for stateful ops"
-            raise ValueError(msg)
-
-        def decorate_stateful(
-            fn: Callable[Concatenate[_S, _P], AsyncGenerator[_T, _S]],
-        ) -> StatefulFn[_P, _S, _T]:
-            return StatefulFn(fn, initial, reducer)
+    if fn is not None:
+        return fn
 
-        return decorate_stateful
+    def decorate(fn: Callable[_P, _T_co]) -> Callable[_P, _T_co]:
+        return fn
 
-    return EffectFn
+    return decorate