diff --git a/burr/core/graph.py b/burr/core/graph.py index fcbc9c5cb..9edc6704e 100644 --- a/burr/core/graph.py +++ b/burr/core/graph.py @@ -57,7 +57,7 @@ def _validate_actions(actions: Optional[List[Action]]): def _validate_transitions( transitions: Optional[List[Tuple[str, str, Condition]]], actions: Set[str] ): - exhausted = {} # items for which we have seen a default transition + exhausted = set() # sources for which we have seen a default transition for from_, to, condition in transitions: if from_ not in actions: raise ValueError( @@ -69,16 +69,28 @@ def _validate_transitions( f"Transition target: `{to}` not found in actions! " f"Please add to actions using with_actions({to}=...)" ) - if condition.name == "default": # we have seen a default transition - if from_ in exhausted: - raise ValueError( - f"Transition `{from_}` -> `{to}` is redundant -- " - f"a default transition has already been set for `{from_}`" - ) - exhausted[from_] = True + _validate_transition_order(from_, to, condition, exhausted) return True +def _validate_transition_order( + from_: str, to: str, condition: Condition, exhausted: Set[str] +) -> None: + """Ensure default transitions are terminal for each source action.""" + if condition.name == "default": + if from_ in exhausted: + raise ValueError( + f"Transition `{from_}` -> `{to}` is redundant -- " + f"a default transition has already been set for `{from_}`" + ) + exhausted.add(from_) + elif from_ in exhausted: + raise ValueError( + f"Transition `{from_}` -> `{to}` is unreachable -- " + f"a default transition has already been set for `{from_}`" + ) + + def _render_graphviz( graphviz_obj, output_file_path: Union[str, pathlib.Path], @@ -328,6 +340,12 @@ def with_transitions( :param transitions: Transitions to add :return: The application builder for future chaining. """ + exhausted = { + from_ + for from_, _, existing_condition in self.transitions + if existing_condition.name == "default" + } + new_transitions = [] for transition in transitions: from_, to_, *conditions = transition if len(conditions) > 0: @@ -341,7 +359,9 @@ def with_transitions( raise ValueError(f"Transition source must be a string, not {action}") if not isinstance(to_, str): raise ValueError(f"Transition target must be a string, not {to_}") - self.transitions.append((action, to_, condition)) + _validate_transition_order(action, to_, condition, exhausted) + new_transitions.append((action, to_, condition)) + self.transitions.extend(new_transitions) return self def with_graph(self, graph: Graph) -> "GraphBuilder": @@ -356,11 +376,22 @@ def with_graph(self, graph: Graph) -> "GraphBuilder": self.actions = [] if self.transitions is None: self.transitions = [] + exhausted = { + from_ + for from_, _, existing_condition in self.transitions + if existing_condition.name == "default" + } + new_transitions = [] + for transition in graph.transitions: + from_, to, condition = ( + transition.from_.name, + transition.to.name, + transition.condition, + ) + _validate_transition_order(from_, to, condition, exhausted) + new_transitions.append((from_, to, condition)) self.actions.extend(graph.actions) - self.transitions.extend( - (transition.from_.name, transition.to.name, transition.condition) - for transition in graph.transitions - ) + self.transitions.extend(new_transitions) return self def build(self) -> Graph: diff --git a/tests/core/test_graph.py b/tests/core/test_graph.py index 22cdda7b6..39921c732 100644 --- a/tests/core/test_graph.py +++ b/tests/core/test_graph.py @@ -95,6 +95,17 @@ def test__validate_transitions_redundant_transition(): ) +def test__validate_transitions_default_must_be_last(): + with pytest.raises(ValueError, match="unreachable"): + _validate_transitions( + [ + ("counter", "result", default), + ("counter", "counter", Condition.expr("count < 10")), + ], + {"counter", "result"}, + ) + + def test__validate_actions_valid(): _validate_actions([Result("test")]) @@ -132,6 +143,12 @@ def test_graph_builder_builds(): assert len(graph.transitions) == 2 +def test_graph_builder_rejects_transition_after_default(): + builder = GraphBuilder().with_transitions(("counter", "result")) + with pytest.raises(ValueError, match="unreachable"): + builder.with_transitions(("counter", "counter", Condition.expr("count < 10"))) + + def test_graph_builder_with_graph(): graph1 = ( GraphBuilder()