From 56a38ee85964b834613854af01e7ade1021d7af2 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 14 Nov 2023 11:17:39 +0000 Subject: [PATCH] Change Channel.checkpoint() to return any python value, add checkpoint tests --- permchain/channels/base.py | 9 +++-- permchain/channels/binop.py | 11 +++--- permchain/channels/context.py | 4 +- permchain/channels/last_value.py | 11 +++--- permchain/channels/topic.py | 19 +++++---- tests/test_channels.py | 68 ++++++++++++++++++++++++++++++++ 6 files changed, 97 insertions(+), 25 deletions(-) diff --git a/permchain/channels/base.py b/permchain/channels/base.py index 7a32a76d4..ee53993f5 100644 --- a/permchain/channels/base.py +++ b/permchain/channels/base.py @@ -15,6 +15,7 @@ from typing_extensions import Self Value = TypeVar("Value") Update = TypeVar("Update") +Checkpoint = TypeVar("Checkpoint") class EmptyChannelError(Exception): @@ -30,7 +31,7 @@ class InvalidUpdateError(Exception): pass -class BaseChannel(Generic[Value, Update], ABC): +class BaseChannel(Generic[Value, Update, Checkpoint], ABC): @property @abstractmethod def ValueType(self) -> Any: @@ -43,7 +44,9 @@ class BaseChannel(Generic[Value, Update], ABC): @contextmanager @abstractmethod - def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: + def empty( + self, checkpoint: Optional[Checkpoint] = None + ) -> Generator[Self, None, None]: """Return a new identical channel, optionally initialized from a checkpoint.""" @asynccontextmanager @@ -68,7 +71,7 @@ class BaseChannel(Generic[Value, Update], ABC): Raises EmptyChannelError if the channel is empty (never updated yet).""" @abstractmethod - def checkpoint(self) -> str | None: + def checkpoint(self) -> Checkpoint | None: """Return a string representation of the channel's current state, or None if the channel doesn't support checkpoints. diff --git a/permchain/channels/binop.py b/permchain/channels/binop.py index 981322b46..322f99c3e 100644 --- a/permchain/channels/binop.py +++ b/permchain/channels/binop.py @@ -1,4 +1,3 @@ -import json from contextlib import contextmanager from typing import Callable, Generator, Generic, Optional, Sequence, Type @@ -7,7 +6,7 @@ from typing_extensions import Self from permchain.channels.base import BaseChannel, EmptyChannelError, Value -class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value]): +class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]): """Stores the result of applying a binary operator to the current value and each new value. ```python @@ -32,10 +31,10 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value]): return self.typ @contextmanager - def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: + def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]: empty = self.__class__(self.typ, self.operator) if checkpoint is not None: - empty.value = json.loads(checkpoint) + empty.value = checkpoint try: yield empty finally: @@ -60,8 +59,8 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value]): except AttributeError: raise EmptyChannelError() - def checkpoint(self) -> str: + def checkpoint(self) -> Value: try: - return json.dumps(self.value) + return self.value except AttributeError: raise EmptyChannelError() diff --git a/permchain/channels/context.py b/permchain/channels/context.py index 183fdafd6..97a95727d 100644 --- a/permchain/channels/context.py +++ b/permchain/channels/context.py @@ -22,7 +22,7 @@ from permchain.channels.base import ( ) -class Context(Generic[Value], BaseChannel[Value, None]): +class Context(Generic[Value], BaseChannel[Value, None, None]): """Exposes the value of a context manager, for the duration of an invocation. Context manager is entered before the first step, and exited after the last step. Optionally, provide an equivalent async context manager, which will be used @@ -66,7 +66,7 @@ class Context(Generic[Value], BaseChannel[Value, None]): raise InvalidUpdateError() @contextmanager - def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: + def empty(self, checkpoint: None = None) -> Generator[Self, None, None]: if self.ctx is None: raise ValueError("Cannot enter sync context manager.") diff --git a/permchain/channels/last_value.py b/permchain/channels/last_value.py index 6dc60e138..33f5cb0b8 100644 --- a/permchain/channels/last_value.py +++ b/permchain/channels/last_value.py @@ -1,4 +1,3 @@ -import json from contextlib import contextmanager from typing import Generator, Generic, Optional, Sequence, Type @@ -12,7 +11,7 @@ from permchain.channels.base import ( ) -class LastValue(Generic[Value], BaseChannel[Value, Value]): +class LastValue(Generic[Value], BaseChannel[Value, Value, Value]): """Stores the last value received, can receive at most one value per step.""" def __init__(self, typ: Type[Value]) -> None: @@ -29,10 +28,10 @@ class LastValue(Generic[Value], BaseChannel[Value, Value]): return self.typ @contextmanager - def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: + def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]: empty = self.__class__(self.typ) if checkpoint is not None: - empty.value = json.loads(checkpoint) + empty.value = checkpoint try: yield empty finally: @@ -55,8 +54,8 @@ class LastValue(Generic[Value], BaseChannel[Value, Value]): except AttributeError: raise EmptyChannelError() - def checkpoint(self) -> str: + def checkpoint(self) -> Value: try: - return json.dumps(self.value) + return self.value except AttributeError: raise EmptyChannelError() diff --git a/permchain/channels/topic.py b/permchain/channels/topic.py index 71b821efe..03907725d 100644 --- a/permchain/channels/topic.py +++ b/permchain/channels/topic.py @@ -1,4 +1,3 @@ -import json from contextlib import contextmanager from typing import Any, Generator, Generic, Iterator, Optional, Sequence, Type, Union @@ -15,7 +14,10 @@ def flatten(values: Sequence[Value | list[Value]]) -> Iterator[Value]: yield value -class Topic(Generic[Value], BaseChannel[Sequence[Value], Value | list[Value]]): +class Topic( + Generic[Value], + BaseChannel[Sequence[Value], Value | list[Value], tuple[set[Value], list[Value]]], +): """A configurable PubSub Topic. Args: @@ -46,12 +48,13 @@ class Topic(Generic[Value], BaseChannel[Sequence[Value], Value | list[Value]]): return Union[self.typ, list[self.typ]] # type: ignore[name-defined] @contextmanager - def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]: + def empty( + self, checkpoint: Optional[tuple[set[Value], list[Value]]] = None + ) -> Generator[Self, None, None]: empty = self.__class__(self.typ, self.unique, self.accumulate) if checkpoint is not None: - parsed = json.loads(checkpoint) - empty.seen = set(parsed["seen"]) - empty.values = list(parsed["values"]) + empty.seen = checkpoint[0] + empty.values = checkpoint[1] try: yield empty finally: @@ -72,5 +75,5 @@ class Topic(Generic[Value], BaseChannel[Sequence[Value], Value | list[Value]]): def get(self) -> Sequence[Value]: return list(self.values) - def checkpoint(self) -> str: - return json.dumps({"seen": list(self.seen), "values": self.values}) + def checkpoint(self) -> tuple[set[Value], list[Value]]: + return (self.seen, self.values) diff --git a/tests/test_channels.py b/tests/test_channels.py index 2b562da62..ed4e28028 100644 --- a/tests/test_channels.py +++ b/tests/test_channels.py @@ -27,6 +27,9 @@ def test_last_value() -> None: assert channel.get() == 3 channel.update([4]) assert channel.get() == 4 + checkpoint = channel.checkpoint() + with LastValue(int).empty(checkpoint) as channel: + assert channel.get() == 4 async def test_last_value_async() -> None: @@ -43,6 +46,9 @@ async def test_last_value_async() -> None: assert channel.get() == 3 channel.update([4]) assert channel.get() == 4 + checkpoint = channel.checkpoint() + async with LastValue(int).aempty(checkpoint) as channel: + assert channel.get() == 4 def test_topic() -> None: @@ -56,6 +62,11 @@ def test_topic() -> None: assert channel.get() == ["c", "d", "d"] channel.update([]) assert channel.get() == [] + channel.update(["e"]) + assert channel.get() == ["e"] + checkpoint = channel.checkpoint() + with Topic(str).empty(checkpoint) as channel: + assert channel.get() == ["e"] async def test_topic_async() -> None: @@ -69,6 +80,11 @@ async def test_topic_async() -> None: assert channel.get() == ["b", "c", "d", "d"] channel.update([]) assert channel.get() == [] + channel.update(["e"]) + assert channel.get() == ["e"] + checkpoint = channel.checkpoint() + async with Topic(str).aempty(checkpoint) as channel: + assert channel.get() == ["e"] def test_topic_unique() -> None: @@ -82,6 +98,13 @@ def test_topic_unique() -> None: assert channel.get() == ["c", "d"], "de-dupes from current and previous steps" channel.update([]) assert channel.get() == [] + channel.update(["e"]) + assert channel.get() == ["e"] + checkpoint = channel.checkpoint() + with Topic(str, unique=True).empty(checkpoint) as channel: + assert channel.get() == ["e"] + channel.update(["d", "f"]) + assert channel.get() == ["f"], "de-dupes from checkpoint" async def test_topic_unique_async() -> None: @@ -95,6 +118,13 @@ async def test_topic_unique_async() -> None: assert channel.get() == ["c", "d"], "de-dupes from current and previous steps" channel.update([]) assert channel.get() == [] + channel.update(["e"]) + assert channel.get() == ["e"] + checkpoint = channel.checkpoint() + async with Topic(str, unique=True).aempty(checkpoint) as channel: + assert channel.get() == ["e"] + channel.update(["d", "f"]) + assert channel.get() == ["f"], "de-dupes from checkpoint" def test_topic_accumulate() -> None: @@ -108,6 +138,11 @@ def test_topic_accumulate() -> None: assert channel.get() == ["a", "b", "b", "c", "d", "d"] channel.update([]) assert channel.get() == ["a", "b", "b", "c", "d", "d"] + checkpoint = channel.checkpoint() + with Topic(str, accumulate=True).empty(checkpoint) as channel: + assert channel.get() == ["a", "b", "b", "c", "d", "d"] + channel.update(["e"]) + assert channel.get() == ["a", "b", "b", "c", "d", "d", "e"] async def test_topic_accumulate_async() -> None: @@ -121,6 +156,11 @@ async def test_topic_accumulate_async() -> None: assert channel.get() == ["a", "b", "b", "c", "d", "d"] channel.update([]) assert channel.get() == ["a", "b", "b", "c", "d", "d"] + checkpoint = channel.checkpoint() + async with Topic(str, accumulate=True).aempty(checkpoint) as channel: + assert channel.get() == ["a", "b", "b", "c", "d", "d"] + channel.update(["e"]) + assert channel.get() == ["a", "b", "b", "c", "d", "d", "e"] def test_topic_unique_accumulate() -> None: @@ -134,6 +174,11 @@ def test_topic_unique_accumulate() -> None: assert channel.get() == ["a", "b", "c", "d"] channel.update([]) assert channel.get() == ["a", "b", "c", "d"] + checkpoint = channel.checkpoint() + with Topic(str, unique=True, accumulate=True).empty(checkpoint) as channel: + assert channel.get() == ["a", "b", "c", "d"] + channel.update(["d", "e"]) + assert channel.get() == ["a", "b", "c", "d", "e"] async def test_topic_unique_accumulate_async() -> None: @@ -147,6 +192,11 @@ async def test_topic_unique_accumulate_async() -> None: assert channel.get() == ["a", "b", "c", "d"] channel.update([]) assert channel.get() == ["a", "b", "c", "d"] + checkpoint = channel.checkpoint() + async with Topic(str, unique=True, accumulate=True).aempty(checkpoint) as channel: + assert channel.get() == ["a", "b", "c", "d"] + channel.update(["d", "e"]) + assert channel.get() == ["a", "b", "c", "d", "e"] def test_binop() -> None: @@ -161,6 +211,9 @@ def test_binop() -> None: assert channel.get() == 6 channel.update([4]) assert channel.get() == 10 + checkpoint = channel.checkpoint() + with BinaryOperatorAggregate(int, operator.add).empty(checkpoint) as channel: + assert channel.get() == 10 async def test_binop_async() -> None: @@ -175,6 +228,9 @@ async def test_binop_async() -> None: assert channel.get() == 6 channel.update([4]) assert channel.get() == 10 + checkpoint = channel.checkpoint() + async with BinaryOperatorAggregate(int, operator.add).aempty(checkpoint) as channel: + assert channel.get() == 10 def test_ctx_manager(mocker: MockerFixture) -> None: @@ -217,6 +273,18 @@ def test_ctx_manager_ctx(mocker: MockerFixture) -> None: with pytest.raises(InvalidUpdateError): channel.update([5]) # type: ignore + checkpoint = channel.checkpoint() + assert checkpoint is None + with Context(httpx.Client).empty(checkpoint) as channel: + assert channel.ValueType is httpx.Client + with pytest.raises(InvalidUpdateError): + assert channel.UpdateType is None + + assert isinstance(channel.get(), httpx.Client) + + with pytest.raises(InvalidUpdateError): + channel.update([5]) + async def test_ctx_manager_async(mocker: MockerFixture) -> None: setup = mocker.Mock()