Add test

permchain/__init__.py

@@ -1,12 +1,10 @@
 from permchain.checkpoint.base import BaseCheckpointAdapter, CheckpointAt
 from permchain.pregel import Channel, Pregel, ReservedChannels
-from permchain.pregel.read import ChannelRead
 
 __all__ = [
     "Channel",
     "Pregel",
     "ReservedChannels",
-    "ChannelRead",
     "BaseCheckpointAdapter",
     "CheckpointAt",
 ]

permchain/pregel/__init__.py

@@ -117,7 +117,7 @@ class Channel:
         """Writes to channels the result of the lambda, or None to skip writing."""
         return ChannelWrite(
             channels=(
-                [(c, RunnablePassthrough()) for c in channels]
+                [(c, None) for c in channels]
                 + [(k, _coerce_write_value(v)) for k, v in kwargs.items()]
             )
         )

permchain/pregel/write.py

@@ -15,7 +15,7 @@ TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None]
 
 
 class ChannelWrite(RunnablePassthrough):
-    channels: Sequence[tuple[str, Runnable]]
+    channels: Sequence[tuple[str, Runnable | None]]
     """
     Mapping of write channels to Runnables that return the value to be written,
     or None to skip writing.
@@ -27,7 +27,7 @@ class ChannelWrite(RunnablePassthrough):
     def __init__(
         self,
         *,
-        channels: Sequence[tuple[str, Runnable]],
+        channels: Sequence[tuple[str, Runnable | None]],
     ):
         super().__init__(func=self._write, afunc=self._awrite, channels=channels)
 
@@ -44,12 +44,17 @@ class ChannelWrite(RunnablePassthrough):
         ]
 
     def _write(self, input: Any, config: RunnableConfig) -> None:
-        values = [(chan, r.invoke(input, config)) for chan, r in self.channels]
+        values = [
+            (chan, r.invoke(input, config) if r else input) for chan, r in self.channels
+        ]
 
         self.do_write(config, **dict(values))
 
     async def _awrite(self, input: Any, config: RunnableConfig) -> None:
-        values = [(chan, await r.ainvoke(input, config)) for chan, r in self.channels]
+        values = [
+            (chan, await r.ainvoke(input, config) if r else input)
+            for chan, r in self.channels
+        ]
 
         self.do_write(config, **dict(values))
 

tests/test_pregel.py

@@ -37,6 +37,7 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
     assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"}
     assert app.output_schema.schema() == {"title": "PregelOutput", "type": "integer"}
     assert app.invoke(2) == 3
+    assert repr(app), "does not raise recursion error"
 
 
 def test_invoke_single_process_in_out_implicit_channels(mocker: MockerFixture) -> None: