wip feat: stream merge

Author: brian14708

src/duron/stream.py

@@ -8,13 +8,15 @@
 from enum import Enum, auto
 from typing import (
     TYPE_CHECKING,
+    Any,
     Concatenate,
     Generic,
     ParamSpec,
     Protocol,
     TypeVar,
     cast,
     final,
+    overload,
 )
 
 from typing_extensions import override
@@ -31,6 +33,7 @@
 _In = TypeVar("_In", contravariant=True)
 _Out = TypeVar("_Out", covariant=True)
 _T = TypeVar("_T")
+_SentialObject = object()
 
 
 class EndOfStream(Exception):
@@ -128,6 +131,25 @@ async def close(self) -> None:
             await self._stop()
             self._state = StreamState.CLOSED
 
+    @overload
+    async def reduce(self, fn: Callable[[_T, _Out], _T], initial: _T, /) -> _T: ...
+    @overload
+    async def reduce(self, fn: Callable[[_Out, _Out], _Out], /) -> _Out: ...
+    async def reduce(
+        self, fn: Callable[[Any, _Out], Any], initial: Any = _SentialObject, /
+    ) -> Any:
+        if initial is _SentialObject:
+            await self.start()
+            acc = await self._get()
+            async for v in self:
+                acc = fn(acc, v)
+            return acc
+        else:
+            acc = initial
+            async for v in self:
+                acc = fn(acc, v)
+            return acc
+
     @final
     def get_nowait(self) -> _Out:
         self._ensure(StreamState.STARTED)
@@ -159,13 +181,16 @@ async def _start(self) -> None: ...
     async def _stop(self) -> None: ...
 
     def map(self, fn: Callable[[_Out], _T]) -> Stream[_T]:
-        self._ensure(StreamState.CONSUMED)
         return _MapStream(self, fn)
 
     def fork(self) -> tuple[Stream[_Out], Stream[_Out]]:
-        self._ensure(StreamState.CONSUMED)
         return _IntoBuffer(self).fork()
 
+    def merge(self, other: Stream[_Out]) -> Stream[_Out]:
+        a = self if isinstance(self, _BufferedStream) else _IntoBuffer(self)
+        b = other if isinstance(other, _BufferedStream) else _IntoBuffer(other)
+        return a.merge(b)
+
 
 @final
 class _AsyncIterableStream(Generic[_Out], Stream[_Out]):
@@ -196,6 +221,7 @@ def _get_nowait(self) -> _Out:
 class _MapStream(Generic[_In, _T], Stream[_T]):
     def __init__(self, source: Stream[_In], fn: Callable[[_In], _T]) -> None:
         super().__init__()
+        source._ensure(StreamState.CONSUMED)
         self._source = source
         self._fn = fn
 
@@ -219,12 +245,13 @@ class _Sentinel:
 
 
 class _BufferedStream(Generic[_T], Stream[_T]):
-    def __init__(self, parent: _BufferedStream[_T] | None = None) -> None:
+    def __init__(self, parents: list[_BufferedStream[_T]] | None = None) -> None:
         super().__init__()
         self.__buffer: deque[_T | _Sentinel] = deque()
         self.__subscribers: list[_BufferedStream[_T]] = []
-        self.__parent = parent
+        self.__parent = parents
         self.__event = asyncio.Event()
+        self.__closed = 1
 
     @override
     async def _get(self) -> _T:
@@ -251,6 +278,9 @@ def _send(self, value: _T):
             s._send(value)
 
     def _send_close(self, exc: BaseException | None = None):
+        self.__closed -= 1
+        if self.__closed != 0:
+            return
         self.__buffer.append(_Sentinel(exc or EndOfStream))
         self.__event.set()
         for s in self.__subscribers:
@@ -260,20 +290,39 @@ def _send_close(self, exc: BaseException | None = None):
     async def _start(self) -> None:
         await super()._start()
         if self.__parent:
-            await self.__parent.start()
+            for p in self.__parent:
+                await p.start()
 
     @override
     def fork(self) -> tuple[Stream[_T], Stream[_T]]:
-        parent = self.__parent or self
+        parent = self.__parent or [self]
         clone: _BufferedStream[_T] = _BufferedStream(parent)
-        parent.__subscribers.append(clone)
+        for p in parent:
+            p.__subscribers.append(clone)
         return self, clone
 
+    @override
+    def merge(self, other: Stream[_T]) -> Stream[_T]:
+        if not isinstance(other, _BufferedStream):
+            return super().merge(other)
+
+        self._ensure(StreamState.CONSUMED)
+        other._ensure(StreamState.CONSUMED)
+        parents: list[_BufferedStream[_T]] = []
+        parents.extend(self.__parent or [self])
+        parents.extend(other.__parent or [other])
+        clone: _BufferedStream[_T] = _BufferedStream(parents)
+        for p in parents:
+            p.__subscribers.append(clone)
+        clone.__closed += len(parents) - 1
+        return clone
+
 
 @final
 class _IntoBuffer(Generic[_T], _BufferedStream[_T]):
     def __init__(self, stream: Stream[_T]) -> None:
         super().__init__()
+        stream._ensure(StreamState.CONSUMED)
         self._stream = stream
         self._task: asyncio.Task[None] | None = None
 
@@ -292,7 +341,7 @@ async def pump() -> None:
     @override
     async def _stop(self) -> None:
         if self._task:
-            self._task.cancel()
+            _ = self._task.cancel()
             with contextlib.suppress(asyncio.CancelledError):
                 await self._task
             self._task = None
@@ -326,12 +375,12 @@ def __init__(
     @override
     def on_next(self, val: _In):
         self._current = self._reducer(self._current, val)
-        self._send(self._current)  # pyright: ignore[reportPrivateUsage]
+        self._send(self._current)
 
     @override
     def on_close(self, exc: BaseException | None):
         self._closed = True if exc is None else exc
-        self._send_close(exc)  # pyright: ignore[reportPrivateUsage]
+        self._send_close(exc)
 
     @override
     async def _start(self) -> None:

tests/test_stream.py

@@ -134,14 +134,15 @@ async def f(s: str) -> AsyncGenerator[str, str]:
 @pytest.mark.asyncio
 async def test_stream_generator():
     @fn()
-    async def activity(_ctx: Context) -> None:
+    async def activity(ctx: Context) -> None:
         async def f() -> AsyncGenerator[int]:
             for i in range(100):
                 yield i
-                await asyncio.sleep(0)
+                await asyncio.sleep(ctx.random().randint(1, 10) * 0.001)
 
         stream = Stream[int].from_iterator(f())
-        m = stream.map(lambda x: x * 2)
+        m = stream.map(lambda x: x + 2000)
+        s3 = Stream[int].from_iterator(f())
 
         with pytest.raises(RuntimeError):
             # cannot iterate twice
@@ -155,7 +156,7 @@ async def f() -> AsyncGenerator[int]:
         # def add(a: int, b: int) -> int:
         #     return a + b
 
-        assert await d.collect() == list(range(0, 200, 2))
+        print(await (s3.merge(d)).collect())
 
     log = MemoryLogStorage()
     async with activity.create_task(log) as t: