Change Channel.checkpoint() to return any python value, add checkpoint tests

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.
 

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()

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.")
 

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()

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)

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()