feat: add effect decorator check loop

Author: brian14708

src/duron/_decorator/effect.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+import asyncio
+import functools
+import inspect
+from typing import TYPE_CHECKING, Any, cast
 from typing_extensions import NamedTuple, ParamSpec, TypeVar, overload
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import AsyncGenerator, Callable, Coroutine
 
 
 _T = TypeVar("_T")
@@ -17,6 +20,61 @@ class Reducer(NamedTuple):
     reducer: Callable[[object, object], object]
 
 
+def _check_loop() -> None:
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        return
+
+    from duron.loop import EventLoop  # noqa: PLC0415
+
+    if isinstance(loop, EventLoop):
+        msg = (
+            "Effects cannot be called from within a duron EventLoop. "
+            "Use 'ctx.run()' to execute effects."
+        )
+        raise RuntimeError(msg)  # noqa: TRY004
+
+
+def _wrap_effect(fn: Callable[_P, Coroutine[Any, Any, _T] | _T]) -> Callable[_P, Any]:
+    if inspect.iscoroutinefunction(fn):
+
+        @functools.wraps(fn)
+        async def async_wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _T:
+            _check_loop()
+            return cast("_T", await fn(*args, **kwargs))
+
+        return async_wrapper
+    if inspect.isasyncgenfunction(fn):
+
+        @functools.wraps(fn)
+        async def async_gen_wrapper(
+            *args: _P.args, **kwargs: _P.kwargs
+        ) -> AsyncGenerator[Any, Any]:
+            _check_loop()
+            gen = fn(*args, **kwargs)
+            try:
+                value = await anext(gen)
+                while True:
+                    try:
+                        sent = yield value
+                        value = await gen.asend(sent)
+                    except GeneratorExit:  # noqa: PERF203
+                        await gen.aclose()
+                        raise
+            except StopAsyncIteration:
+                return
+
+        return async_gen_wrapper
+
+    @functools.wraps(fn)
+    def sync_wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _T:
+        _check_loop()
+        return cast("_T", fn(*args, **kwargs))
+
+    return sync_wrapper
+
+
 @overload
 def effect(fn: Callable[_P, _T], /) -> Callable[_P, _T]: ...
 @overload
@@ -50,9 +108,5 @@ async def counter(
 
     """
     if fn is not None:
-        return fn
-
-    def decorate(fn: Callable[_P, _T]) -> Callable[_P, _T]:
-        return fn
-
-    return decorate
+        return _wrap_effect(fn)
+    return _wrap_effect

tests/test_session.py

@@ -9,13 +9,14 @@
 
 import pytest
 
-from duron import Context, Provided, Session, Stream, StreamWriter, durable
+from duron import Context, Provided, Session, Stream, StreamWriter, durable, effect
 from duron.contrib.codecs import PickleCodec
 from duron.contrib.storage import MemoryLogStorage
 
 
 @pytest.mark.asyncio
 async def test_invoke() -> None:
+    @effect
     async def u() -> str:
         for _ in range(random.randint(1, 10)):
             await asyncio.sleep(0.001)