Get new tests passing

Author: dmontagu

pydantic_graph/pydantic_graph/beta/decision.py

@@ -178,20 +178,19 @@ def to(
 
     def fork(
         self,
-        get_forks: Callable[[Self], Sequence[Decision[StateT, DepsT, HandledT | SourceT]]],
+        get_forks: Callable[[Self], Sequence[DecisionBranch[SourceT]]],
         /,
     ) -> DecisionBranch[SourceT]:
         """Create a fork in the execution path.
 
         Args:
-            get_forks: Function that generates fork decisions.
+            get_forks: Function that generates forked decision branches.
 
         Returns:
             A completed DecisionBranch with forked execution paths.
         """
-        n_initial_branches = len(self.decision.branches)
-        fork_decisions = get_forks(self)
-        new_paths = [b.path for fd in fork_decisions for b in fd.branches[n_initial_branches:]]
+        fork_decision_branches = get_forks(self)
+        new_paths = [b.path for b in fork_decision_branches]
         return DecisionBranch(source=self.source, matches=self.matches, path=self.path_builder.fork(new_paths))
 
     def transform(

pydantic_graph/pydantic_graph/beta/graph.py

@@ -693,14 +693,39 @@ def _handle_path(self, path: Path, inputs: Any, fork_stack: ForkStack) -> Sequen
 
     def _handle_edges(self, node: AnyNode, inputs: Any, fork_stack: ForkStack) -> Sequence[GraphTask]:
         edges = self.graph.edges_by_source.get(node.id, [])
-        assert len(edges) == 1 or isinstance(node, Fork), (
+        assert len(edges) == 1 or (isinstance(node, Fork) and not node.is_map), (
             edges,
             node.id,
         )  # this should have already been ensured during graph building
 
         new_tasks: list[GraphTask] = []
-        for path in edges:
-            new_tasks.extend(self._handle_path(path, inputs, fork_stack))
+
+        if isinstance(node, Fork):
+            node_run_id = NodeRunID(str(uuid.uuid4()))
+            if node.is_map:
+                # Eagerly raise a clear error if the input value is not iterable as expected
+                try:
+                    iter(inputs)
+                except TypeError:
+                    raise RuntimeError(f'Cannot map non-iterable value: {inputs!r}')
+
+                # If the map specifies a downstream join id, eagerly create a reducer for it
+                if (join_id := node.downstream_join_id) is not None:
+                    join_node = self.graph.nodes[join_id]
+                    assert isinstance(join_node, Join)
+                    self._active_reducers[(join_id, node_run_id)] = join_node.create_reducer(), fork_stack
+
+                for thread_index, input_item in enumerate(inputs):
+                    item_tasks = self._handle_path(
+                        edges[0], input_item, fork_stack + (ForkStackItem(node.id, node_run_id, thread_index),)
+                    )
+                    new_tasks += item_tasks
+            else:
+                for i, path in enumerate(edges):
+                    new_tasks += self._handle_path(path, inputs, fork_stack + (ForkStackItem(node.id, node_run_id, i),))
+        else:
+            new_tasks += self._handle_path(edges[0], inputs, fork_stack)
+
         return new_tasks
 
     def _is_fork_run_completed(self, tasks: Iterable[GraphTask], join_id: JoinID, fork_run_id: NodeRunID) -> bool:

pydantic_graph/pydantic_graph/beta/graph_builder.py

@@ -341,7 +341,7 @@ def join(
             return join(reducer_type=reducer_factory, node_id=node_id)
 
     # Edge building
-    def add(self, *edges: EdgePath[StateT, DepsT]) -> None:
+    def add(self, *edges: EdgePath[StateT, DepsT]) -> None:  # noqa C901
         """Add one or more edge paths to the graph.
 
         This method processes edge paths and automatically creates any necessary
@@ -359,12 +359,12 @@ def _handle_path(p: Path):
             """
             for item in p.items:
                 if isinstance(item, BroadcastMarker):
-                    new_node = Fork[Any, Any](id=item.fork_id, is_map=False)
+                    new_node = Fork[Any, Any](id=item.fork_id, is_map=False, downstream_join_id=None)
                     self._insert_node(new_node)
                     for path in item.paths:
                         _handle_path(Path(items=[*path.items]))
                 elif isinstance(item, MapMarker):
-                    new_node = Fork[Any, Any](id=item.fork_id, is_map=True)
+                    new_node = Fork[Any, Any](id=item.fork_id, is_map=True, downstream_join_id=item.downstream_join_id)
                     self._insert_node(new_node)
                 elif isinstance(item, DestinationMarker):
                     pass
@@ -710,6 +710,7 @@ def build(self) -> Graph[StateT, DepsT, GraphInputT, GraphOutputT]:
         # TODO(P3): Consider doing a deepcopy here to prevent modifications to the underlying nodes and edges
         nodes = self._nodes
         edges_by_source = self._edges_by_source
+        nodes, edges_by_source = _flatten_paths(nodes, edges_by_source)
         nodes, edges_by_source = _normalize_forks(nodes, edges_by_source)
         parent_forks = _collect_dominating_forks(nodes, edges_by_source)
 
@@ -726,6 +727,52 @@ def build(self) -> Graph[StateT, DepsT, GraphInputT, GraphOutputT]:
         )
 
 
+def _flatten_paths(
+    nodes: dict[NodeID, AnyNode], edges: dict[NodeID, list[Path]]
+) -> tuple[dict[NodeID, AnyNode], dict[NodeID, list[Path]]]:
+    new_nodes = nodes.copy()
+    new_edges: dict[NodeID, list[Path]] = defaultdict(list)
+
+    paths_to_handle: list[tuple[NodeID, Path]] = []
+
+    def _split_at_first_fork(path: Path) -> tuple[Path, list[tuple[NodeID, Path]]]:
+        for i, item in enumerate(path.items):
+            if isinstance(item, MapMarker):
+                if item.fork_id not in nodes:
+                    new_nodes[item.fork_id] = Fork(
+                        id=item.fork_id, is_map=True, downstream_join_id=item.downstream_join_id
+                    )
+                upstream = Path(list(path.items[:i]) + [DestinationMarker(item.fork_id)])
+                downstream = Path(path.items[i + 1 :])
+                return upstream, [(item.fork_id, downstream)]
+
+            if isinstance(item, BroadcastMarker):
+                if item.fork_id not in nodes:
+                    new_nodes[item.fork_id] = Fork(id=item.fork_id, is_map=True, downstream_join_id=None)
+                upstream = Path(list(path.items[:i]) + [DestinationMarker(item.fork_id)])
+                return upstream, [(item.fork_id, p) for p in item.paths]
+        return path, []
+
+    for node in new_nodes.values():
+        if isinstance(node, Decision):
+            for branch in node.branches:
+                upstream, downstreams = _split_at_first_fork(branch.path)
+                branch.path = upstream
+                paths_to_handle.extend(downstreams)
+
+    for source_id, edges_from_source in edges.items():
+        for path in edges_from_source:
+            paths_to_handle.append((source_id, path))
+
+    while paths_to_handle:
+        source_id, path = paths_to_handle.pop()
+        upstream, downstreams = _split_at_first_fork(path)
+        new_edges[source_id].append(upstream)
+        paths_to_handle.extend(downstreams)
+
+    return new_nodes, dict(new_edges)
+
+
 def _normalize_forks(
     nodes: dict[NodeID, AnyNode], edges: dict[NodeID, list[Path]]
 ) -> tuple[dict[NodeID, AnyNode], dict[NodeID, list[Path]]]:
@@ -756,25 +803,11 @@ def _normalize_forks(
         if len(edges_from_source) == 1:
             new_edges[source_id] = edges_from_source
             continue
-        new_fork = Fork[Any, Any](id=ForkID(NodeID(f'{node.id}_broadcast_fork')), is_map=False)
+        new_fork = Fork[Any, Any](id=ForkID(NodeID(f'{node.id}_broadcast_fork')), is_map=False, downstream_join_id=None)
         new_nodes[new_fork.id] = new_fork
         new_edges[source_id] = [Path(items=[BroadcastMarker(fork_id=new_fork.id, paths=edges_from_source)])]
         new_edges[new_fork.id] = edges_from_source
 
-    while paths_to_handle:
-        path = paths_to_handle.pop()
-        for item in path.items:
-            if isinstance(item, MapMarker):
-                assert item.fork_id in new_nodes
-                new_edges[item.fork_id] = [path.next_path]
-                paths_to_handle.append(path.next_path)
-                break
-            elif isinstance(item, BroadcastMarker):
-                assert item.fork_id in new_nodes
-                new_edges[item.fork_id] = [*item.paths]
-                paths_to_handle.extend(item.paths)
-                break
-
     return new_nodes, new_edges
 
 
@@ -808,7 +841,6 @@ def _collect_dominating_forks(
 
         if isinstance(node, Fork):
             fork_ids.add(node.id)
-            continue
 
         def _handle_path(path: Path, last_source_id: NodeID):
             """Process a path and collect edges and fork information.

pydantic_graph/pydantic_graph/beta/join.py

@@ -196,6 +196,8 @@ def finalize(self, ctx: StepContext[object, object, None]) -> dict[K, V]:
 
 
 class SupportsSum(Protocol):
+    """A protocol for a type that supports adding to itself."""
+
     @abstractmethod
     def __add__(self, other: Self, /) -> Self:
         pass

pydantic_graph/pydantic_graph/beta/node.py

@@ -11,7 +11,7 @@
 
 from typing_extensions import TypeVar
 
-from pydantic_graph.beta.id_types import ForkID, NodeID
+from pydantic_graph.beta.id_types import ForkID, JoinID, NodeID
 
 StateT = TypeVar('StateT', infer_variance=True)
 """Type variable for graph state."""
@@ -76,6 +76,8 @@ class Fork(Generic[InputT, OutputT]):
     If True, InputT must be Sequence[OutputT] and each element is sent to a separate branch.
     If False, InputT must be OutputT and the same data is sent to all branches.
     """
+    downstream_join_id: JoinID | None
+    """Optional identifier of a downstream join node that should be jumped to if mapping an empty iterable."""
 
     def _force_variance(self, inputs: InputT) -> OutputT:  # pragma: no cover
         """Force type variance for proper generic typing.

pydantic_graph/pydantic_graph/beta/paths.py

@@ -135,6 +135,7 @@ class Path:
 
     items: Sequence[PathItem]
     """The sequence of path items that define this path."""
+    # TODO: Change items to be Sequence[TransformMarker | MapMarker | LabelMarker] and add field `destination: BroadcastMarker | DestinationMarker`
 
     @property
     def last_fork(self) -> BroadcastMarker | MapMarker | None:

tests/graph/beta/test_decisions.py

@@ -437,7 +437,7 @@ async def handle_b(ctx: StepContext[DecisionState, None, object]) -> str:
 
 async def test_decision_branch_fork():
     """Test DecisionBranchBuilder.fork method."""
-    g = GraphBuilder(state_type=DecisionState, output_type=str)
+    g = GraphBuilder(state_type=DecisionState, output_type=list[str])
 
     @g.step
     async def choose_option(ctx: StepContext[DecisionState, None, None]) -> Literal['fork']:
@@ -453,28 +453,22 @@ async def path_2(ctx: StepContext[DecisionState, None, object]) -> str:
 
     collect = g.join(ListAppendReducer[str])
 
-    @g.step
-    async def combine(ctx: StepContext[DecisionState, None, list[str]]) -> str:
-        return ', '.join(ctx.inputs)
-
     g.add(
         g.edge_from(g.start_node).to(choose_option),
         g.edge_from(choose_option).to(
             g.decision().branch(
                 g.match(TypeExpression[Literal['fork']]).fork(
                     lambda b: [
-                        b.decision.branch(g.match(TypeExpression[Literal['fork']]).to(path_1)),
-                        b.decision.branch(g.match(TypeExpression[Literal['fork']]).to(path_2)),
+                        b.to(path_1),
+                        b.to(path_2),
                     ]
                 )
             )
         ),
         g.edge_from(path_1, path_2).to(collect),
-        g.edge_from(collect).to(combine),
-        g.edge_from(combine).to(g.end_node),
+        g.edge_from(collect).to(g.end_node),
     )
 
     graph = g.build()
     result = await graph.run(state=DecisionState())
-    assert 'Path 1' in result
-    assert 'Path 2' in result
+    assert sorted(result) == ['Path 1', 'Path 2']

tests/graph/beta/test_parent_forks.py

@@ -1,5 +1,7 @@
 """Tests for parent fork identification and dominator analysis."""
 
+from inline_snapshot import snapshot
+
 from pydantic_graph.beta.parent_forks import ParentForkFinder
 
 
@@ -50,7 +52,10 @@ def test_parent_fork_with_cycle():
 
 
 def test_parent_fork_nested_forks():
-    """Test parent fork identification with nested forks."""
+    """Test parent fork identification with nested forks.
+
+    In this case, it should return the most ancestral valid parent fork.
+    """
     join_id = 'J'
     nodes = {'start', 'F1', 'F2', 'A', 'B', 'C', 'J', 'end'}
     start_ids = {'start'}
@@ -68,36 +73,43 @@ def test_parent_fork_nested_forks():
     parent_fork = finder.find_parent_fork(join_id)
 
     assert parent_fork is not None
-    # Should find F2 as the immediate parent fork
-    assert parent_fork.fork_id == 'F2'
+    # Should find F1 as the most ancestral parent fork
+    assert parent_fork.fork_id == 'F1'
 
 
-def test_parent_fork_most_ancestral():
-    """Test that the most ancestral valid parent fork is found."""
-    join_id = 'J'
-    nodes = {'start', 'F1', 'F2', 'I', 'A', 'B', 'C', 'J', 'end'}
+def test_parent_fork_parallel_nested_forks():
+    """Test parent fork identification with nested forks.
+
+    This test is mostly included to document the current behavior, which is always to use the most ancestral
+    valid fork, even if the most ancestral fork isn't guaranteed to pass through the specified join, and another
+    fork is.
+
+    We might want to change this behavior at some point, but if we do, we'll probably want to do so in some sort
+    of user-specified way to ensure we don't break user code.
+    """
+    nodes = {'start', 'F1', 'F2-A', 'F2-B', 'A1', 'A2', 'B1', 'B2', 'C', 'J-A', 'J-B', 'J', 'end'}
     start_ids = {'start'}
-    fork_ids = {'F1', 'F2'}
-    # F1 is the most ancestral fork, F2 is nested, with intermediate node I, and a cycle from J back to I
+    fork_ids = {'F1', 'F2A', 'F2B'}
     edges = {
         'start': ['F1'],
-        'F1': ['F2'],
-        'F2': ['I'],
-        'I': ['A', 'B'],
-        'A': ['J'],
-        'B': ['J'],
-        'J': ['C'],
-        'C': ['end', 'I'],  # Cycle back to I
+        'F1': ['F2-A', 'F2-B'],
+        'F2-A': ['A1', 'A2'],
+        'F2-B': ['B1', 'B2'],
+        'A1': ['J-A'],
+        'A2': ['J-A'],
+        'B1': ['J-B'],
+        'B2': ['J-B'],
+        'J-A': ['J'],
+        'J-B': ['J'],
+        'J': ['end'],
     }
 
     finder = ParentForkFinder(nodes, start_ids, fork_ids, edges)
-    parent_fork = finder.find_parent_fork(join_id)
-
-    # F2 is not a valid parent because J has a cycle back to I which avoids F2
-    # F1 is also not valid for the same reason
-    # But we should find I as the intermediate fork... wait, I is not a fork
-    # So we should get None OR the most ancestral fork that doesn't have the cycle issue
-    assert parent_fork is None or parent_fork.fork_id in fork_ids
+    parent_fork_ids = [
+        finder.find_parent_fork(join_id).fork_id  # pyright: ignore[reportOptionalMemberAccess]
+        for join_id in ['J-A', 'J-B', 'J']
+    ]
+    assert parent_fork_ids == snapshot(['F1', 'F1', 'F1'])  # NOT: ['F2-A', 'F2-B', 'F1'] as one might suspect
 
 
 def test_parent_fork_no_forks():

tests/graph/beta/test_util.py

@@ -1,7 +1,5 @@
 """Tests for pydantic_graph.beta.util module."""
 
-from typing import Union
-
 from pydantic_graph.beta.util import (
     Some,
     TypeExpression,
@@ -18,9 +16,9 @@ def test_type_expression_unpacking():
     assert result is int
 
     # Test with TypeExpression wrapper
-    wrapped = TypeExpression[Union[str, int]]
+    wrapped = TypeExpression[str | int]
     result = unpack_type_expression(wrapped)
-    assert result == Union[str, int]
+    assert result == str | int
 
 
 def test_some_wrapper():