Add more channel types

README.md

@@ -37,6 +37,7 @@ Check `examples` for more examples.
 - [x] Implement input_schema and output_schema in Pregel
 - [ ] More tests
   - [ ] Test different input and output types (str, str sequence, None)
+  - [ ] Add tests for Stream, UniqueInbox
 - [ ] Implement checkpointing
   - [ ] Save checkpoints at end of each step
   - [ ] Load checkpoint at start of invocation

permchain/channels.py

@@ -2,6 +2,7 @@ import json
 from abc import ABC, abstractmethod
 from contextlib import asynccontextmanager, contextmanager
 from typing import (
+    Any,
     AsyncContextManager,
     AsyncGenerator,
     Callable,
@@ -10,8 +11,11 @@ from typing import (
     Generic,
     Optional,
     Sequence,
+    Tuple,
     Type,
     TypeVar,
+    Union,
+    cast,
 )
 from typing import ContextManager as ContextManagerType
 
@@ -32,12 +36,12 @@ class InvalidUpdateError(Exception):
 class Channel(Generic[Value, Update], ABC):
     @property
     @abstractmethod
-    def ValueType(self) -> Type[Value]:
+    def ValueType(self) -> Any:
         """The type of the value stored in the channel."""
 
     @property
     @abstractmethod
-    def UpdateType(self) -> Type[Update]:
+    def UpdateType(self) -> Any:
         """The type of the update received by the channel."""
 
     @contextmanager
@@ -166,7 +170,7 @@ class LastValue(Generic[Value], Channel[Value, Value]):
         return json.dumps(self.value)
 
 
-class Inbox(Generic[Value], Channel[Sequence[Value], Value]):
+class Inbox(Generic[Value], Channel[Sequence[Value], Value | Sequence[Value]]):
     """Stores all values received, resets in each step."""
 
     def __init__(self, typ: Type[Value]) -> None:
@@ -178,9 +182,9 @@ class Inbox(Generic[Value], Channel[Sequence[Value], Value]):
         return Sequence[self.typ]  # type: ignore[name-defined]
 
     @property
-    def UpdateType(self) -> Type[Value]:
+    def UpdateType(self) -> Any:
         """The type of the update received by the channel."""
-        return self.typ
+        return Union[self.typ, Sequence[self.typ]]  # type: ignore[name-defined]
 
     @contextmanager
     def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]:
@@ -195,8 +199,60 @@ class Inbox(Generic[Value], Channel[Sequence[Value], Value]):
             except AttributeError:
                 pass
 
-    def update(self, values: Sequence[Value]) -> None:
-        self.queue = tuple(values)
+    def update(self, values: Sequence[Value | Sequence[Value]]) -> None:
+        self.queue = tuple(
+            cast(Value, v)
+            for value in values
+            for v in ((value,) if isinstance(value, self.typ) else value)
+        )
+
+    def get(self) -> Sequence[Value]:
+        try:
+            return self.queue
+        except AttributeError:
+            raise EmptyChannelError()
+
+    def checkpoint(self) -> str:
+        return json.dumps(self.queue)
+
+
+class UniqueInbox(Generic[Value], Channel[Sequence[Value], Value | Sequence[Value]]):
+    """Stores all unique values received, resets in each step."""
+
+    def __init__(self, typ: Type[Value]) -> None:
+        self.typ = typ
+
+    @property
+    def ValueType(self) -> Type[Sequence[Value]]:
+        """The type of the value stored in the channel."""
+        return Sequence[self.typ]  # type: ignore[name-defined]
+
+    @property
+    def UpdateType(self) -> Any:
+        """The type of the update received by the channel."""
+        return Union[self.typ, Sequence[self.typ]]  # type: ignore[name-defined]
+
+    @contextmanager
+    def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]:
+        empty = self.__class__(self.typ)
+        if checkpoint is not None:
+            empty.queue = tuple(json.loads(checkpoint))
+        try:
+            yield empty
+        finally:
+            try:
+                del empty.queue
+            except AttributeError:
+                pass
+
+    def update(self, values: Sequence[Value | Sequence[Value]]) -> None:
+        self.queue = tuple(
+            set(
+                cast(Value, v)
+                for value in values
+                for v in ((value,) if isinstance(value, self.typ) else value)
+            )
+        )
 
     def get(self) -> Sequence[Value]:
         try:
@@ -213,6 +269,7 @@ class Set(Generic[Value], Channel[FrozenSet[Value], Value]):
 
     def __init__(self, typ: Type[Value]) -> None:
         self.typ = typ
+        self.set = set[Value]()
 
     @property
     def ValueType(self) -> Type[FrozenSet[Value]]:
@@ -232,14 +289,9 @@ class Set(Generic[Value], Channel[FrozenSet[Value], Value]):
         try:
             yield empty
         finally:
-            try:
-                del empty.set
-            except AttributeError:
-                pass
+            pass
 
     def update(self, values: Sequence[Value]) -> None:
-        if not hasattr(self, "set"):
-            self.set = set[Value]()
         self.set.update(values)
 
     def get(self) -> FrozenSet[Value]:
@@ -252,6 +304,46 @@ class Set(Generic[Value], Channel[FrozenSet[Value], Value]):
         return json.dumps(list(self.set))
 
 
+class Stream(Generic[Value], Channel[Tuple[Value], Value]):
+    """Stores all unique values received."""
+
+    def __init__(self, typ: Type[Value]) -> None:
+        self.typ = typ
+        self.set = list[Value]()
+
+    @property
+    def ValueType(self) -> Type[Tuple[Value]]:
+        """The type of the value stored in the channel."""
+        return Tuple[self.typ]  # type: ignore[name-defined]
+
+    @property
+    def UpdateType(self) -> Type[Value]:
+        """The type of the update received by the channel."""
+        return self.typ
+
+    @contextmanager
+    def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]:
+        empty = self.__class__(self.typ)
+        if checkpoint is not None:
+            empty.set = json.loads(checkpoint)
+        try:
+            yield empty
+        finally:
+            pass
+
+    def update(self, values: Sequence[Value]) -> None:
+        self.set.extend(values)
+
+    def get(self) -> Tuple[Value]:
+        try:
+            return tuple(self.set)
+        except AttributeError:
+            raise EmptyChannelError()
+
+    def checkpoint(self) -> str:
+        return json.dumps(self.set)
+
+
 AsyncValue = TypeVar("AsyncValue")
 
 

tests/test_channels.py

@@ -1,6 +1,6 @@
 import operator
 from contextlib import asynccontextmanager, contextmanager
-from typing import AsyncGenerator, FrozenSet, Generator, Sequence
+from typing import AsyncGenerator, FrozenSet, Generator, Sequence, Union
 
 import pytest
 from pytest_mock import MockerFixture
@@ -43,21 +43,21 @@ async def test_last_value_async() -> None:
 def test_inbox() -> None:
     with channels.Inbox(str).empty() as channel:
         assert channel.ValueType is Sequence[str]
-        assert channel.UpdateType is str
+        assert channel.UpdateType is Union[str, Sequence[str]]
 
         with pytest.raises(channels.EmptyChannelError):
             channel.get()
 
         channel.update(["a", "b"])
         assert channel.get() == ("a", "b")
-        channel.update(["c"])
-        assert channel.get() == ("c",)
+        channel.update([["c"], "d"])
+        assert channel.get() == ("c", "d")
 
 
 async def test_inbox_async() -> None:
     async with channels.Inbox(str).aempty() as channel:
         assert channel.ValueType is Sequence[str]
-        assert channel.UpdateType is str
+        assert channel.UpdateType is Union[str, Sequence[str]]
 
         with pytest.raises(channels.EmptyChannelError):
             channel.get()
@@ -65,7 +65,8 @@ async def test_inbox_async() -> None:
         channel.update(["a", "b"])
         assert channel.get() == ("a", "b")
         channel.update(["c"])
-        assert channel.get() == ("c",)
+        channel.update([["c"], "d"])
+        assert channel.get() == ("c", "d")
 
 
 def test_set() -> None:
@@ -73,9 +74,7 @@ def test_set() -> None:
         assert channel.ValueType is FrozenSet[str]
         assert channel.UpdateType is str
 
-        with pytest.raises(channels.EmptyChannelError):
-            channel.get()
-
+        assert channel.get() == frozenset()
         channel.update(["a", "b"])
         assert channel.get() == frozenset(("a", "b"))
         channel.update(["b", "c"])
@@ -87,9 +86,7 @@ async def test_set_async() -> None:
         assert channel.ValueType is FrozenSet[str]
         assert channel.UpdateType is str
 
-        with pytest.raises(channels.EmptyChannelError):
-            channel.get()
-
+        assert channel.get() == frozenset()
         channel.update(["a", "b"])
         assert channel.get() == frozenset(("a", "b"))
         channel.update(["b", "c"])
@@ -147,7 +144,7 @@ def test_ctx_manager(mocker: MockerFixture) -> None:
         assert channel.get() == 5
 
         with pytest.raises(channels.InvalidUpdateError):
-            channel.update([5])
+            channel.update([5])  # type: ignore
 
     assert setup.call_count == 1
     assert cleanup.call_count == 1
@@ -183,7 +180,7 @@ async def test_ctx_manager_async(mocker: MockerFixture) -> None:
         assert channel.get() == 5
 
         with pytest.raises(channels.InvalidUpdateError):
-            channel.update([5])
+            channel.update([5])  # type: ignore
 
     assert setup.call_count == 1
     assert cleanup.call_count == 1