wip stream writer context manager

Author: brian14708

src/duron/_core/invoke.py

@@ -41,9 +41,7 @@
 
 if TYPE_CHECKING:
     from asyncio.exceptions import CancelledError
-    from collections.abc import Callable, Coroutine
-    from contextvars import Token
-    from types import TracebackType
+    from collections.abc import AsyncGenerator, Callable, Coroutine
 
     from duron._core.ops import Op, StreamObserver
     from duron._decorator.durable import DurableFn
@@ -63,16 +61,17 @@
     from duron.typing import FunctionType
 
 
-_T_co = TypeVar("_T_co", covariant=True)
 _T = TypeVar("_T")
 _P = ParamSpec("_P")
+_OutT = TypeVar("_OutT", covariant=True)  # noqa: PLC0105
 
 _CURRENT_VERSION: Final = 0
 
 
-def invoke(
-    fn: DurableFn[_P, _T_co], log: LogStorage, /, *, tracer: Tracer | None = None
-) -> contextlib.AbstractAsyncContextManager[DurableRun[_P, _T_co]]:
+@contextlib.asynccontextmanager
+async def invoke(
+    fn: DurableFn[_P, _OutT], log: LogStorage, /, *, tracer: Tracer | None = None
+) -> AsyncGenerator[DurableRun[_P, _OutT], None]:
     """Create an invocation context for this durable function.
 
     Args:
@@ -82,18 +81,24 @@ def invoke(
 
     Returns:
         Async context manager for Invoke instance
-    """
-    return _InvokeGuard(DurableRun(fn, log), tracer)
+    """  # noqa: DOC402, DOC202
+    token = current_tracer.set(tracer)
+    run = DurableRun(fn, log)
+    try:
+        yield run
+    finally:
+        await run.close()
+        current_tracer.reset(token)
 
 
 @final
-class DurableRun(Generic[_P, _T_co]):
+class DurableRun(Generic[_P, _OutT]):
     __slots__ = ("_fn", "_log", "_run", "_watchers")
 
-    def __init__(self, fn: DurableFn[_P, _T_co], log: LogStorage) -> None:
+    def __init__(self, fn: DurableFn[_P, _OutT], log: LogStorage) -> None:
         self._fn = fn
         self._log = log
-        self._run: _InvokeRun | None = None
+        self._run: _DurableRun | None = None
         self._watchers: list[tuple[dict[str, str], StreamObserver]] = []
 
     async def start(self, *args: _P.args, **kwargs: _P.kwargs) -> None:
@@ -111,20 +116,20 @@ async def prelude() -> InitParams:  # noqa: RUF029
         type_info = inspect_function(self._fn.fn)
         p = _invoke_prelude(self._fn, type_info, prelude)
         loop = await create_loop()
-        self._run = _InvokeRun(loop, p, self._log, codec, watchers=self._watchers)
+        self._run = _DurableRun(loop, p, self._log, codec, watchers=self._watchers)
         await self._run.resume()
 
     async def resume(self) -> None:
         """Resume a previously started invocation."""
         type_info = inspect_function(self._fn.fn)
         prelude = _invoke_prelude(self._fn, type_info, _resume_init)
         loop = await create_loop()
-        self._run = _InvokeRun(
+        self._run = _DurableRun(
             loop, prelude, self._log, self._fn.codec, watchers=self._watchers
         )
         await self._run.resume()
 
-    async def wait(self) -> _T_co:
+    async def wait(self) -> _OutT:
         """Wait for the durable function invocation to complete \
                 and return its result.
 
@@ -137,7 +142,7 @@ async def wait(self) -> _T_co:
         if self._run is None:
             msg = "Job not started"
             raise RuntimeError(msg)
-        return cast("_T_co", await self._run.run())
+        return cast("_OutT", await self._run.run())
 
     async def close(self) -> None:
         if self._run:
@@ -214,37 +219,13 @@ def open_stream(
             return stream
         return _StreamWriter(self, name)
 
-    def get_run(self) -> _InvokeRun:
+    def get_run(self) -> _DurableRun:
         if self._run is None:
             msg = "Job not started"
             raise RuntimeError(msg)
         return self._run
 
 
-@final
-class _InvokeGuard(Generic[_P, _T_co]):
-    __slots__ = ("_job", "_token", "_tracer")
-
-    def __init__(self, job: DurableRun[_P, _T_co], tracer: Tracer | None) -> None:
-        self._job = job
-        self._tracer = tracer
-        self._token: Token[Tracer | None] | None = None
-
-    async def __aenter__(self) -> DurableRun[_P, _T_co]:
-        self._token = current_tracer.set(self._tracer)
-        return self._job
-
-    async def __aexit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_value: BaseException | None,
-        traceback: TracebackType | None,
-    ) -> None:
-        await self._job.close()
-        if self._token:
-            current_tracer.reset(self._token)
-
-
 class InitParams(TypedDict):
     version: int
     args: list[JSONValue]
@@ -253,10 +234,10 @@ class InitParams(TypedDict):
 
 
 async def _invoke_prelude(
-    job_fn: DurableFn[..., _T_co],
+    job_fn: DurableFn[..., _OutT],
     type_info: FunctionType,
     init: Callable[[], Coroutine[Any, Any, InitParams]],
-) -> _T_co:
+) -> _OutT:
     loop = asyncio.get_running_loop()
     assert isinstance(loop, EventLoop)  # noqa: S101
 
@@ -316,7 +297,7 @@ async def _invoke_prelude(
 
 
 @final
-class _InvokeRun:
+class _DurableRun:
     __slots__ = (
         "_codec",
         "_lease",