feat wip stream

Author: brian14708

src/duron/context.py

@@ -1,24 +1,29 @@
 from __future__ import annotations
 
+import asyncio
 from contextvars import ContextVar
 from random import Random
-from typing import TYPE_CHECKING, ParamSpec, TypeVar, cast, final
+from typing import TYPE_CHECKING, Concatenate, ParamSpec, TypeVar, cast, final
 
 from typing_extensions import overload
 
-from duron.ops import FnCall
+from duron.ops import FnCall, StreamCreate
+from duron.stream import StreamTask
 
 if TYPE_CHECKING:
-    from collections.abc import Awaitable, Callable
+    from collections.abc import AsyncGenerator, Awaitable, Callable
     from contextvars import Token
     from types import TracebackType
 
     from duron.event_loop import EventLoop
     from duron.fn import Fn
+    from duron.stream import Observer, Stream
 
     _T = TypeVar("_T")
+    _S = TypeVar("_S")
     _P = ParamSpec("_P")
 
+
 _context: ContextVar[Context | None] = ContextVar("duron_context", default=None)
 
 
@@ -74,6 +79,8 @@ async def run(
         *args: _P.args,
         **kwargs: _P.kwargs,
     ) -> _T:
+        if asyncio.get_event_loop() is not self._loop:
+            raise RuntimeError("Context time can only be used in the context loop")
         type_info = self._task.codec.inspect_function(fn)
         return cast(
             "_T",
@@ -87,11 +94,46 @@ async def run(
             ),
         )
 
+    def run_stream(
+        self,
+        initial: _T,
+        reducer: Callable[[_T, _S], _T],
+        fn: Callable[Concatenate[_T, _P], AsyncGenerator[_S, _T]],
+        /,
+        *args: _P.args,
+        **kwargs: _P.kwargs,
+    ) -> StreamTask[_S, _T]:
+        if asyncio.get_event_loop() is not self._loop:
+            raise RuntimeError("Context time can only be used in the context loop")
+        return StreamTask(
+            self._loop,
+            initial,
+            reducer,
+            fn,
+            *args,
+            **kwargs,
+        )
+
+    async def create_stream(self, observer: Observer[_T] | None) -> Stream[_T]:
+        if asyncio.get_event_loop() is not self._loop:
+            raise RuntimeError("Context time can only be used in the context loop")
+        o = cast("Observer[object]", observer) if observer else None
+        return cast(
+            "Stream[_T]",
+            await self._loop.create_op(StreamCreate(observer=o)),
+        )
+
     def time(self) -> float:
+        if asyncio.get_event_loop() is not self._loop:
+            raise RuntimeError("Context time can only be used in the context loop")
         return self._loop.time()
 
     def time_ns(self) -> int:
+        if asyncio.get_event_loop() is not self._loop:
+            raise RuntimeError("Context time can only be used in the context loop")
         return self._loop.time_ns()
 
     def random(self) -> Random:
+        if asyncio.get_event_loop() is not self._loop:
+            raise RuntimeError("Context random can only be used in the context loop")
         return Random(self._loop.generate_op_id())

src/duron/event_loop.py

@@ -5,6 +5,7 @@
 import contextvars
 import heapq
 import logging
+import os
 from asyncio import AbstractEventLoop, Handle, Task, TimerHandle, events
 from collections import deque
 from dataclasses import dataclass
@@ -47,7 +48,7 @@ class OpFuture(asyncio.Future[object]):
     id: bytes
     params: object
 
-    def __init__(self, id: bytes, params: object, loop: EventLoop) -> None:
+    def __init__(self, id: bytes, params: object, loop: AbstractEventLoop) -> None:
         super().__init__(loop=loop)
         self.id = id
         self.params = params
@@ -212,11 +213,36 @@ def poll_completion(self, task: Future[_T]) -> WaitSet | None:
         finally:
             events._set_running_loop(old)
 
-    def create_op(self, params: object) -> OpFuture:
-        id = self.generate_op_id()
-        s = OpFuture(id, params, self)
-        self._ops[id] = s
-        return s
+    def create_op(
+        self, params: object, /, *, loop: AbstractEventLoop | None = None
+    ) -> OpFuture:
+        if loop is not None and loop is not self:
+            id = os.urandom(12)
+            host_fut: asyncio.Future[object] = OpFuture(id, params, loop)
+            op_fut = OpFuture(id, params, self)
+            self._ops[id] = op_fut
+
+            def op_to_host(f: asyncio.Future[object]):
+                if f.cancelled():
+                    _ = loop.call_soon(host_fut.cancel)
+                elif e := f.exception():
+                    _ = loop.call_soon(host_fut.set_exception, e)
+                else:
+                    _ = loop.call_soon(host_fut.set_result, f.result())
+
+            def host_to_op(f: asyncio.Future[object]):
+                if f.cancelled() and not op_fut.done():
+                    _ = self.call_soon(op_fut.cancel)
+
+            op_fut.add_done_callback(op_to_host)
+            host_fut.add_done_callback(host_to_op)
+            self._event.set()
+            return host_fut
+        else:
+            id = self.generate_op_id()
+            op_fut = OpFuture(id, params, self)
+            self._ops[id] = op_fut
+            return op_fut
 
     @overload
     def post_completion(

src/duron/log.py

@@ -46,8 +46,7 @@ class StreamCreateEntry(_BaseEntry):
 class StreamEmitEntry(_BaseEntry):
     type: Literal["stream/emit"]
     stream_id: str
-    value: NotRequired[JSONValue]
-    state: NotRequired[JSONValue]
+    value: JSONValue
 
 
 class StreamCompleteEntry(_BaseEntry):

src/duron/ops.py

@@ -6,6 +6,8 @@
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable, Coroutine
 
+    from duron.stream import Observer
+
 
 @dataclass(slots=True)
 class FnCall:
@@ -21,4 +23,21 @@ class TaskRun:
     return_type: type | None = None
 
 
-Op = FnCall | TaskRun
+@dataclass(slots=True)
+class StreamCreate:
+    observer: Observer[object] | None
+
+
+@dataclass(slots=True)
+class StreamEmit:
+    stream_id: str
+    value: object
+
+
+@dataclass(slots=True)
+class StreamClose:
+    stream_id: str
+    exception: BaseException | None
+
+
+Op = FnCall | StreamCreate | StreamEmit | StreamClose | TaskRun

src/duron/stream.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+import asyncio
+from typing import (
+    TYPE_CHECKING,
+    Concatenate,
+    Generic,
+    ParamSpec,
+    Protocol,
+    TypeVar,
+    cast,
+    final,
+)
+
+from typing_extensions import override
+
+from duron.ops import FnCall, StreamClose, StreamCreate, StreamEmit
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, Callable
+
+    from duron.event_loop import EventLoop
+
+    _P = ParamSpec("_P")
+
+_In = TypeVar("_In", contravariant=True)
+_Out = TypeVar("_Out", covariant=True)
+
+
+class Observer(Generic[_In], Protocol):
+    def on_next(self, value: _In, /) -> None: ...
+    def on_close(self, error: BaseException | None, /) -> None: ...
+
+
+@final
+class Stream(Generic[_In]):
+    def __init__(self, id: str, loop: EventLoop) -> None:
+        self._stream_id = id
+        self._loop = loop
+        self._event = asyncio.Event()
+
+    async def send(self, value: _In, /) -> None:
+        _ = await self._loop.create_op(
+            StreamEmit(stream_id=self._stream_id, value=value),
+            loop=asyncio.get_running_loop(),
+        )
+
+    async def close(self, error: BaseException | None = None, /) -> None:
+        _ = await self._loop.create_op(
+            StreamClose(stream_id=self._stream_id, exception=error),
+            loop=asyncio.get_running_loop(),
+        )
+        self._event.set()
+
+    async def wait(self) -> None:
+        _ = await self._event.wait()
+
+
+@final
+class _StreamObserver(Generic[_In, _Out], Observer[_In]):
+    def __init__(self, initial: _Out, reducer: Callable[[_Out, _In], _Out]):
+        self.current = initial
+        self.enable = True
+        self._reducer = reducer
+        self.data: list[_Out] = []
+        self.closed: bool | BaseException = False
+
+    @override
+    def on_next(self, val: _In):
+        if self.enable:
+            self.current = self._reducer(self.current, val)
+            self.data.append(self.current)
+
+    @override
+    def on_close(self, exc: BaseException | None):
+        self.closed = True if exc is None else exc
+
+
+@final
+class StreamTask(Generic[_In, _Out]):
+    def __init__(
+        self,
+        loop: EventLoop,
+        initial: _Out,
+        reducer: Callable[[_Out, _In], _Out],
+        fn: Callable[Concatenate[_Out, _P], AsyncGenerator[_In, _Out]],
+        /,
+        *args: _P.args,
+        **kwargs: _P.kwargs,
+    ) -> None:
+        self._loop = loop
+        self._reducer = reducer
+        self._fn = fn
+        self._args = args
+        self._kwargs = kwargs
+        self._obs = _StreamObserver(initial, self._reducer)
+        self._op = self._loop.create_op(
+            StreamCreate(observer=cast("Observer[object]", self._obs))
+        )
+        self._queue: asyncio.Queue[tuple[_Out] | None | BaseException] = (
+            self._setup_stream()
+        )
+
+    def __aiter__(self) -> StreamTask[_In, _Out]:
+        return self
+
+    async def __anext__(self) -> _Out:
+        item = await self._queue.get()
+        if item is None:
+            raise StopAsyncIteration
+        if isinstance(item, BaseException):
+            raise item
+        return item[0]
+
+    async def discard(self) -> None:
+        async for _ in self:
+            ...
+
+    def _setup_stream(self) -> asyncio.Queue[tuple[_Out] | None | BaseException]:
+        queue: asyncio.Queue[tuple[_Out] | None | BaseException] = asyncio.Queue()
+
+        async def worker():
+            stream = cast("Stream[_In]", await self._op)
+            try:
+                state = self._obs.current
+                self._obs.enable = False
+
+                for d in self._obs.data:
+                    await queue.put((d,))
+                if self._obs.closed is True:
+                    return
+                elif isinstance(self._obs.closed, BaseException):
+                    raise self._obs.closed
+
+                gen = self._fn(state, *self._args, **self._kwargs)
+                state_partial = await gen.__anext__()
+
+                while True:
+                    state = self._reducer(state, state_partial)
+                    await stream.send(state_partial)
+                    await queue.put((state,))
+                    state_partial = await gen.asend(state)
+            except StopAsyncIteration as _e:
+                await stream.close()
+            except BaseException as e:
+                await queue.put(e)
+                raise
+            finally:
+                await queue.put(None)
+
+        _ = self._loop.create_op(
+            FnCall(callable=worker, args=(), kwargs={}, return_type=None)
+        )
+        return queue