route() returns Moar

src/wf_authoring/builder/__init__.py

@@ -1,4 +1,4 @@
 from .core import WorkflowBuilder
-from .refs import BranchRef, StepRef
+from .refs import BranchRef, RouteRef, StepRef
 
-__all__ = ["BranchRef", "StepRef", "WorkflowBuilder"]
+__all__ = ["BranchRef", "RouteRef", "StepRef", "WorkflowBuilder"]

src/wf_authoring/builder/core.py

@@ -36,7 +36,7 @@ from .mapping import (
     coerce_path,
     normalize_mapping,
 )
-from .refs import BranchRef, StepRef, is_node_spec, step_id
+from .refs import BranchRef, RouteRef, StepRef, is_node_spec, step_id
 
 
 @dataclass(slots=True)
@@ -239,7 +239,7 @@ class WorkflowBuilder:
         *,
         id: str | None = None,
         default: BranchRef = runtime_error,
-    ) -> dict[object, StepRef]:
+    ) -> RouteRef:
         """Route graph data by equality checks or one boolean condition.
 
         `branch()` wires outcomes already produced by a node. `route()` is the
@@ -259,9 +259,10 @@ class WorkflowBuilder:
         *,
         id: str | None,
         default: BranchRef,
-    ) -> dict[object, StepRef]:
+    ) -> RouteRef:
         """Expand value cases into an ordered chain of equality checks."""
         resolved_targets: dict[object, StepRef] = {}
+        conditions: list[ConditionNode] = []
         default_target = self.use(default) if is_node_spec(default) else default
         previous_condition: ConditionNode | None = None
         condition_base = id or "condition"
@@ -270,6 +271,7 @@ class WorkflowBuilder:
                 id=self._next_step_id(condition_base),
                 check=value.eq(case_value),
             )
+            conditions.append(condition)
             resolved = self.use(target) if is_node_spec(target) else target
             if previous_condition is not None:
                 self.connect(previous_condition, "false", condition)
@@ -280,7 +282,11 @@ class WorkflowBuilder:
             raise ValueError("WorkflowBuilder.route requires at least one case")
         self.connect(previous_condition, "false", cast(StepRef, default_target))
         resolved_targets["default"] = cast(StepRef, default_target)
-        return resolved_targets
+        return RouteRef(
+            entry=conditions[0],
+            conditions=tuple(conditions),
+            targets=resolved_targets,
+        )
 
     def _route_condition(
         self,
@@ -289,7 +295,7 @@ class WorkflowBuilder:
         *,
         id: str | None,
         default: BranchRef,
-    ) -> dict[object, StepRef]:
+    ) -> RouteRef:
         """Route a boolean condition expression through true/false outcomes."""
         invalid_cases = [case for case in cases if not isinstance(case, bool)]
         if invalid_cases:
@@ -304,7 +310,11 @@ class WorkflowBuilder:
             resolved = self.use(target) if is_node_spec(target) else target
             self.connect(condition_node, outcome, cast(StepRef, resolved))
             resolved_targets[case_value] = cast(StepRef, resolved)
-        return resolved_targets
+        return RouteRef(
+            entry=condition_node,
+            conditions=(condition_node,),
+            targets=resolved_targets,
+        )
 
     def compile(self) -> Workflow:
         if self.start is None:

src/wf_authoring/builder/refs.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from dataclasses import dataclass
 from typing import Any, TypeAlias, TypeGuard
 
 from wf_core import ConditionNode, ForeachNode, InterruptNode, NodeUse
@@ -10,6 +11,19 @@ StepRef: TypeAlias = str | NodeUse | ConditionNode | ForeachNode | InterruptNode
 BranchRef: TypeAlias = StepRef | NodeSpec[Any, Any]
 
 
+@dataclass(frozen=True, slots=True)
+class RouteRef:
+    """Reference bundle returned by route() for generated condition nodes."""
+
+    entry: ConditionNode
+    conditions: tuple[ConditionNode, ...]
+    targets: dict[object, StepRef]
+
+    def __getitem__(self, key: object) -> StepRef:
+        """Keep route["case"] ergonomic while exposing generated conditions."""
+        return self.targets[key]
+
+
 def step_id(ref: StepRef) -> str:
     """Return the core step id for either a step object or an id string."""
     if isinstance(ref, str):

tests/authoring/test_builder_branch.py

@@ -83,10 +83,11 @@ def test_builder_route_expands_state_value_cases_into_condition_chain() -> None:
         default=fallback,
     )
 
-    assert [node.id for node in builder.nodes if node.type == "condition"] == [
+    assert [node.id for node in targets.conditions] == [
         "condition",
         "condition_2",
     ]
+    assert targets.entry.id == "condition"
     assert [(edge.from_, edge.outcome, edge.to) for edge in builder.edges] == [
         ("condition", "true", "left"),
         ("condition", "false", "condition_2"),
@@ -134,6 +135,8 @@ def test_builder_route_accepts_boolean_condition_expression() -> None:
 
     targets = builder.route(state("count").ge(1), {True: left, False: right})
 
+    assert targets.entry.id == "condition"
+    assert [node.id for node in targets.conditions] == ["condition"]
     assert [(edge.from_, edge.outcome, edge.to) for edge in builder.edges] == [
         ("condition", "true", "left"),
         ("condition", "false", "right"),

tests/rewrite/workflow.py

@@ -6,7 +6,6 @@ from tests.rewrite.actions import (
     post_roll_router,
     pre_roll_router,
     prep,
-    rate_booster,
     roll,
     tick,
 )
@@ -48,7 +47,7 @@ gacha.connect("tick", "ok", "counter_up")
 gacha.connect("counter_up", "ok", "rate_booster")
 # gacha.connect("rate_booster", "0", rate_same)
 # gacha.connect("rate_booster", "65", rate_up)
-gacha.route(
+rate_route = gacha.route(
     state("counter.c_80").ge(65),
     {
         True: rate_up,
@@ -56,6 +55,7 @@ gacha.route(
     },
     id = "rate_booster"
 )
+
 gacha.use(pre_roll_router, id="router")
 
 gacha.connect("rate_up", "ok", "router")