set_next

Author: davesh0812

storey/flow.py

@@ -77,6 +77,7 @@ def __init__(
             self.name = type(self).__name__
 
         self._closeables = []
+        self._selected_outlet: Optional[list[str]] = None
 
     def _init(self):
         self._termination_received = 0
@@ -277,10 +278,14 @@ async def _do_downstream(self, event, outlets=None):
         if outlets:
             outlets = outlets
         elif event is not _termination_obj:
-            if asyncio.iscoroutinefunction(self.select_outlets):
-                outlet_names = await self.select_outlets(event.body)
+            if self._selected_outlet:
+                outlet_names = self._selected_outlet
+                self._selected_outlet = None
             else:
-                outlet_names = self.select_outlets(event.body)
+                if asyncio.iscoroutinefunction(self.select_outlets):
+                    outlet_names = await self.select_outlets(event.body)
+                else:
+                    outlet_names = self.select_outlets(event.body)
             outlets = self._check_outlets_by_names(outlet_names) if outlet_names else self._outlets
         else:
             outlets = self._outlets
@@ -403,6 +408,15 @@ def select_outlets(self, event) -> typing.Optional[Collection[str]]:
         """
         return None
 
+    def set_next_outlets(self, outlet_names: Union[str, list[str]]):
+        """
+        Set the next outlets to which the event will be sent. This method can be used in conjunction with
+        select_outlets() to dynamically determine the outlets for the next event.
+
+        :param outlet_names: A collection of outlet names to which the next event should be sent.
+        """
+        self._selected_outlet = outlet_names if isinstance(outlet_names, list) else [outlet_names]
+
     def _check_outlets_by_names(self, outlet_names: Collection[str]) -> list["Flow"]:
         outlets = []
 
@@ -448,7 +462,7 @@ class Choice(Flow):
     def _init(self):
         super()._init()
         # TODO: hacky way of supporting mlrun preview, which replaces targets with a DFTarget
-        self._passthrough_for_preview = list(self._name_to_outlet) == ["dataframe"] if self._name_to_outlet else False
+        self._passthrough_for_preview = list(self._name_to_outlet) == ["dataframe"]
 
     async def _do(self, event):
         if event is _termination_obj:

tests/test_flow.py

@@ -5480,7 +5480,8 @@ def select_outlets(self, event):
         return outlets
 
 
-def test_regular_step_with_choice():
+@pytest.mark.parametrize("fn_select_outlets", [True, False])
+def test_regular_step_with_choice(fn_select_outlets):
     class MyStep(Map):
         def select_outlets(self, event):
             outlets = ["all_events"]
@@ -5490,7 +5491,60 @@ def select_outlets(self, event):
                 outlets.append("up_to_five")
             return outlets
 
+    def select(event):
+        outlets = ["all_events"]
+        if event > 5:
+            outlets.append("more_than_five")
+        else:
+            outlets.append("up_to_five")
+        return outlets
+
+    source = SyncEmitSource()
+    if fn_select_outlets:
+        my_step = Map(fn=lambda x: x, termination_result_fn=lambda x, y: x + y, fn_select_outlets=select)
+    else:
+        my_step = MyStep(fn=lambda x: x, termination_result_fn=lambda x, y: x + y)
+
+    all_events = Map(lambda x: x, name="all_events")
+    more_than_five = Map(lambda x: x * 10, name="more_than_five")
+    up_to_five = Map(lambda x: x * 100, name="up_to_five")
+    sum_up_all_events = Reduce(0, lambda acc, x: acc + x)
+    sum_up_more_than_five = Reduce(0, lambda acc, x: acc + x)
+    sum_up_up_to_five = Reduce(0, lambda acc, x: acc + x)
+
+    source.to(my_step)
+    my_step.to(all_events)
+    my_step.to(more_than_five)
+    my_step.to(up_to_five)
+    all_events.to(sum_up_all_events)
+    more_than_five.to(sum_up_more_than_five)
+    up_to_five.to(sum_up_up_to_five)
+
+    controller = source.run()
+
+    for i in range(4, 8):
+        controller.emit(i)
+
+    controller.terminate()
+    termination_result = controller.await_termination()
+
+    expected = sum(range(4, 8)) + sum(range(6, 8)) * 10 + sum(range(4, 6)) * 100
+    assert termination_result == expected
+
+
+def test_regular_step_with_set_next():
+    class MyStep(MapClass):
+        def do(self, event):
+            outlets = ["all_events"]
+            if event > 5:
+                outlets.append("more_than_five")
+            else:
+                outlets.append("up_to_five")
+            self.set_next_outlets(outlets)
+            return event
+
     source = SyncEmitSource()
+
     my_step = MyStep(fn=lambda x: x, termination_result_fn=lambda x, y: x + y)
     all_events = Map(lambda x: x, name="all_events")
     more_than_five = Map(lambda x: x * 10, name="more_than_five")