diff --git a/permchain/pubsub.py b/permchain/pubsub.py index 2b3ddb068..087a80e8d 100644 --- a/permchain/pubsub.py +++ b/permchain/pubsub.py @@ -24,20 +24,28 @@ T = TypeVar("T") T_in = TypeVar("T_in") T_out = TypeVar("T_out") +Process = RunnableSubscriber[T_in] | RunnableReducer[T_in] + class PubSub(Runnable[Any, Any], ABC): - processes: Sequence[RunnableSubscriber[Any] | RunnableReducer[Any]] + processes: Sequence[Process] connection: PubSubConnection def __init__( self, - processes: Sequence[RunnableSubscriber[Any] | RunnableReducer[Any]], + *procs: Process | Sequence[Process], + processes: Sequence[Process] = (), connection: PubSubConnection, ) -> None: super().__init__() - self.processes = processes self.connection = connection + self.processes = list(processes) + for proc in procs: + if isinstance(proc, Sequence): + self.processes.extend(proc) + else: + self.processes.append(proc) def with_retry(self, **kwargs: Any) -> Runnable[Any, Any]: return self.__class__( diff --git a/tests/test_invoke.py b/tests/test_invoke.py index 022034453..8f36e385d 100644 --- a/tests/test_invoke.py +++ b/tests/test_invoke.py @@ -16,7 +16,7 @@ def test_invoke_single_process_in_out(mocker: MockerFixture): assert chain.invoke(2) == 3 conn = InMemoryPubSubConnection() - pubsub = PubSub(processes=(chain,), connection=conn) + pubsub = PubSub(chain, connection=conn) # Using in-memory conn internals to make assertions about pubsub # If we start with 0 listeners @@ -38,7 +38,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture): assert chain_two.invoke(2) == 3 conn = InMemoryPubSubConnection() - pubsub = PubSub(processes=(chain_one, chain_two), connection=conn) + pubsub = PubSub(chain_one, chain_two, connection=conn) # Using in-memory conn internals to make assertions about pubsub # If we start with 0 listeners @@ -60,8 +60,8 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture): assert chain_two.invoke(2) == 3 conn = InMemoryPubSubConnection(clear_on_disconnect=False) - pubsub_one = PubSub(processes=(chain_one,), connection=conn) - pubsub_two = PubSub(processes=(chain_two,), connection=conn) + pubsub_one = PubSub(chain_one, connection=conn) + pubsub_two = PubSub(chain_two, connection=conn) # Using in-memory conn internals to make assertions about pubsub # If we start with 0 listeners