diff --git a/README.md b/README.md index 8e02dbba6..ace1c193f 100644 --- a/README.md +++ b/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 diff --git a/permchain/channels.py b/permchain/channels.py index 71ccdf3a4..12e5e08a7 100644 --- a/permchain/channels.py +++ b/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") diff --git a/tests/test_channels.py b/tests/test_channels.py index 2d95e49cd..1ebbc8693 100644 --- a/tests/test_channels.py +++ b/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