Add more coverage

Author: dmontagu

pydantic_graph/pydantic_graph/beta/decision.py

@@ -61,9 +61,6 @@ def branch(self, branch: DecisionBranch[T]) -> Decision[StateT, DepsT, HandledT
 
         Returns:
             A new Decision with the additional branch.
-
-        Note:
-            TODO(P3): Add an overload that skips the need for `match`, and is just less flexible about the building.
         """
         return Decision(id=self.id, branches=self.branches + [branch], note=self.note)
 

pydantic_graph/pydantic_graph/beta/graph.py

@@ -256,7 +256,7 @@ def render(self, *, title: str | None = None, direction: StateDiagramDirection |
         """
         from pydantic_graph.beta.mermaid import build_mermaid_graph
 
-        return build_mermaid_graph(self).render(title=title, direction=direction)
+        return build_mermaid_graph(self.nodes, self.edges_by_source).render(title=title, direction=direction)
 
     def __repr__(self) -> str:
         super_repr = super().__repr__()  # include class and memory address
@@ -550,15 +550,15 @@ def _handle_result(result: EndMarker[OutputT] | JoinItem | Sequence[GraphTask])
                         len(afs) > len(fork_stack) and fork_stack == afs[: len(fork_stack)]
                         for afs in active_fork_stacks
                     ):
-                        assert False  # TODO: Need to cover this in a test
+                        # TODO: Need to cover this in a test
                         continue  # this join_state is a strict prefix for one of the other active join_states
                     self._active_reducers.pop((join_id, fork_run_id))  # we're handling it now, so we can pop it
                     join_node = self.graph.nodes[join_id]
                     assert isinstance(join_node, Join), f'Expected a `Join` but got {join_node}'
                     new_tasks = self._handle_edges(join_node, join_state.current, join_state.downstream_fork_stack)
                     maybe_overridden_result = yield new_tasks  # give an opportunity to override these
                     if _handle_result(maybe_overridden_result):
-                        assert False  # TODO: Need to cover this in a test
+                        # TODO: Need to cover this in a test
                         return
 
         raise RuntimeError(  # pragma: no cover
@@ -665,36 +665,11 @@ def _handle_path(self, path: Path, inputs: Any, fork_stack: ForkStack) -> Sequen
             return []  # pragma: no cover
 
         item = path.items[0]
+        assert not isinstance(item, MapMarker | BroadcastMarker), (
+            'These markers should be removed from paths during graph building'
+        )
         if isinstance(item, DestinationMarker):
             return [GraphTask(item.destination_id, inputs, fork_stack)]
-        elif isinstance(item, MapMarker):
-            # Eagerly raise a clear error if the input value is not iterable as expected
-            try:
-                assert False  # TODO: Need to cover this in a test
-                iter(inputs)
-            except TypeError:
-                raise RuntimeError(f'Cannot map non-iterable value: {inputs!r}')
-
-            node_run_id = NodeRunID(str(uuid.uuid4()))
-
-            # If the map specifies a downstream join id, eagerly create a join state for it
-            if item.downstream_join_id is not None:
-                join_node = self.graph.nodes[item.downstream_join_id]
-                assert isinstance(join_node, Join)
-                self._active_reducers[(item.downstream_join_id, node_run_id)] = JoinState(
-                    join_node.initial_factory(), fork_stack
-                )
-
-            map_tasks: list[GraphTask] = []
-            for thread_index, input_item in enumerate(inputs):
-                item_tasks = self._handle_path(
-                    path.next_path, input_item, fork_stack + (ForkStackItem(item.fork_id, node_run_id, thread_index),)
-                )
-                map_tasks += item_tasks
-            return map_tasks
-        elif isinstance(item, BroadcastMarker):
-            assert False  # TODO: Need to cover this in a test
-            return [GraphTask(item.fork_id, inputs, fork_stack)]
         elif isinstance(item, TransformMarker):
             inputs = item.transform(StepContext(state=self.state, deps=self.deps, inputs=inputs))
             return self._handle_path(path.next_path, inputs, fork_stack)

pydantic_graph/pydantic_graph/beta/graph_builder.py

@@ -22,6 +22,7 @@
 from pydantic_graph.beta.graph import Graph
 from pydantic_graph.beta.id_types import ForkID, JoinID, NodeID, generate_placeholder_node_id, replace_placeholder_id
 from pydantic_graph.beta.join import Join, JoinNode, ReducerFunction
+from pydantic_graph.beta.mermaid import build_mermaid_graph
 from pydantic_graph.beta.node import (
     EndNode,
     Fork,
@@ -59,6 +60,12 @@
 T = TypeVar('T', infer_variance=True)
 
 
+class GraphBuildingError(ValueError):
+    """An error raised during graph-building."""
+
+    pass
+
+
 @dataclass(init=False)
 class GraphBuilder(Generic[StateT, DepsT, GraphInputT, GraphOutputT]):
     """A builder for constructing executable graph definitions.
@@ -476,7 +483,7 @@ def match_node(
         Returns:
             A DecisionBranch for the BaseNode type
         """
-        assert False  # TODO: Need to cover this in a test
+        # TODO: Need to cover this in a test
         node = NodeStep(source)
         path = Path(items=[DestinationMarker(node.id)])
         return DecisionBranch(source=source, matches=matches, path=path, destinations=[node])
@@ -532,7 +539,9 @@ def _insert_node(self, node: AnyNode) -> None:
         elif isinstance(existing, NodeStep) and isinstance(node, NodeStep) and existing.node_type is node.node_type:
             pass
         elif existing is not node:
-            raise ValueError(f'All nodes must have unique node IDs. {node.id!r} was the ID for {existing} and {node}')
+            raise GraphBuildingError(
+                f'All nodes must have unique node IDs. {node.id!r} was the ID for {existing} and {node}'
+            )
 
     def _edge_from_return_hint(
         self, node: SourceNode[StateT, DepsT, Any], return_hint: TypeOrTypeExpression[Any]
@@ -631,9 +640,6 @@ def build(self) -> Graph[StateT, DepsT, GraphInputT, GraphOutputT]:
         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)
-        print(nodes)
-        print(edges_by_source)
-        print(parent_forks)
 
         return Graph[StateT, DepsT, GraphInputT, GraphOutputT](
             name=self.name,
@@ -720,10 +726,9 @@ def _normalize_forks(
         if len(edges_from_source) == 1:
             new_edges[source_id] = edges_from_source
             continue
-        assert False  # TODO: Need to cover this in a test, specifically by using `.to()`  with multiple arguments
         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[source_id] = [Path(items=[DestinationMarker(new_fork.id)])]
         new_edges[new_fork.id] = edges_from_source
 
     return new_nodes, new_edges
@@ -796,9 +801,22 @@ def _handle_path(path: Path, last_source_id: NodeID):
             join.id, explicit_fork_id=join.parent_fork_id, prefer_closest=join.preferred_parent_fork == 'closest'
         )
         if dominating_fork is None:
-            # TODO(P3): Print out the mermaid graph and explain the problem
-            assert False  # TODO: Need to cover this in a test
-            raise ValueError(f'Join node {join.id} has no dominating fork')
+            rendered_mermaid_graph = build_mermaid_graph(graph_nodes, graph_edges_by_source).render()
+            error_message = f"""\
+For every Join J in the graph, there must be a Fork F between the StartNode and J satisfying:
+* Every path from the StartNode to J passes through F
+* There are no cycles in the graph including both J and F.
+In this case, F is called a "dominating fork" for J.
+
+This is used to determine when all tasks upstream of this Join are complete and we can proceed with execution.
+
+Mermaid diagram:
+{rendered_mermaid_graph}
+
+Join {join.id!r} in this graph has no dominating fork.\
+"""
+            # TODO: Need to cover this in a test
+            raise GraphBuildingError(error_message)
         dominating_forks[join.id] = dominating_fork
 
     return dominating_forks

pydantic_graph/pydantic_graph/beta/mermaid.py

@@ -2,15 +2,15 @@
 
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Any, Literal
+from typing import Literal
 
 from typing_extensions import assert_never
 
 from pydantic_graph.beta.decision import Decision
-from pydantic_graph.beta.graph import Graph
 from pydantic_graph.beta.id_types import NodeID
 from pydantic_graph.beta.join import Join
 from pydantic_graph.beta.node import EndNode, Fork, StartNode
+from pydantic_graph.beta.node_types import AnyNode
 from pydantic_graph.beta.paths import BroadcastMarker, DestinationMarker, LabelMarker, MapMarker, Path
 from pydantic_graph.beta.step import NodeStep, Step
 
@@ -49,7 +49,9 @@ class MermaidEdge:
     label: str | None
 
 
-def build_mermaid_graph(graph: Graph[Any, Any, Any, Any]) -> MermaidGraph:  # noqa C901
+def build_mermaid_graph(  # noqa C901
+    graph_nodes: dict[NodeID, AnyNode], graph_edges_by_source: dict[NodeID, list[Path]]
+) -> MermaidGraph:
     """Build a mermaid graph."""
     nodes: list[MermaidNode] = []
     edges_by_source: dict[str, list[MermaidEdge]] = defaultdict(list)
@@ -63,7 +65,7 @@ def _collect_edges(path: Path, last_source_id: NodeID) -> None:
             elif isinstance(item, DestinationMarker):
                 edges_by_source[last_source_id].append(MermaidEdge(last_source_id, item.destination_id, working_label))
 
-    for node_id, node in graph.nodes.items():
+    for node_id, node in graph_nodes.items():
         kind: NodeKind
         label: str | None = None
         note: str | None = None
@@ -89,11 +91,11 @@ def _collect_edges(path: Path, last_source_id: NodeID) -> None:
         source_node = MermaidNode(id=node_id, kind=kind, label=label, note=note)
         nodes.append(source_node)
 
-    for k, v in graph.edges_by_source.items():
+    for k, v in graph_edges_by_source.items():
         for path in v:
             _collect_edges(path, k)
 
-    for node in graph.nodes.values():
+    for node in graph_nodes.values():
         if isinstance(node, Decision):
             for branch in node.branches:
                 _collect_edges(branch.path, node.id)

pydantic_graph/pydantic_graph/beta/parent_forks.py

@@ -13,8 +13,6 @@
 In most typical graphs, such dominating forks exist naturally. However, when there are multiple
 subsequent forks, the choice of parent fork can be ambiguous and may need to be specified by
 the graph designer.
-
-TODO(P3): Expand this documentation with more detailed examples and edge cases.
 """
 
 from __future__ import annotations
@@ -69,7 +67,6 @@ class ParentForkFinder(Generic[T]):
     edges: dict[T, list[T]]  # source_id to list of destination_ids
     """Graph edges represented as adjacency list mapping source nodes to destinations."""
 
-    # TODO: Add unit tests of this class that make use of explicit_fork_id and prefer_closest
     def find_parent_fork(
         self, join_id: T, *, explicit_fork_id: T | None = None, prefer_closest: bool = False
     ) -> ParentFork[T] | None:

pydantic_graph/pydantic_graph/beta/paths.py

@@ -44,7 +44,7 @@ class TransformFunction(Protocol[StateT, DepsT, InputT, OutputT]):
         OutputT: The type of the output data
     """
 
-    # TODO: Rework to better-support lambdas through callable union like we did with ReducerFunction
+    # TODO: Consider reworking to better-support lambdas through callable union like we did with ReducerFunction
     def __call__(self, ctx: StepContext[StateT, DepsT, InputT]) -> OutputT:
         """Execute the step function with the given context.
 

pydantic_graph/pydantic_graph/beta/step.py

@@ -22,7 +22,7 @@
 OutputT = TypeVar('OutputT', infer_variance=True)
 
 
-# TODO: Remove inputs from StepContext and provide multiple allowed signatures like with ReducerFunction
+# TODO: Consider removing inputs from StepContext and provide multiple allowed signatures like with ReducerFunction
 @dataclass(init=False)
 class StepContext(Generic[StateT, DepsT, InputT]):
     """Context information passed to step functions during graph execution.

tests/graph/beta/test_graph_builder.py

@@ -7,6 +7,7 @@
 import pytest
 
 from pydantic_graph.beta import GraphBuilder, StepContext
+from pydantic_graph.beta.graph_builder import GraphBuildingError
 from pydantic_graph.beta.join import reduce_list_append, reduce_sum
 from pydantic_graph.beta.node import Fork
 
@@ -262,7 +263,7 @@ async def step_one(ctx: StepContext[SimpleState, None, None]) -> int:
     async def step_two(ctx: StepContext[SimpleState, None, None]) -> int:
         return 2
 
-    with pytest.raises(ValueError, match='All nodes must have unique node IDs'):
+    with pytest.raises(GraphBuildingError, match='All nodes must have unique node IDs'):
         g.add(
             g.edge_from(g.start_node).to(step_one),
             g.edge_from(g.start_node).to(step_two),

tests/graph/beta/test_graph_execution.py

@@ -19,7 +19,7 @@ class ExecutionState:
 
 
 async def test_map_to_end_node_cancels_pending():
-    """Test that mapping directly to end_node cancels pending tasks (covers graph.py:480)."""
+    """Test that mapping directly to end_node cancels pending tasks"""
     import asyncio
 
     g = GraphBuilder(state_type=ExecutionState, output_type=int)
@@ -73,7 +73,7 @@ async def process_item(ctx: StepContext[ExecutionState, None, int]) -> int:
 
 
 async def test_broadcast_marker_handling():
-    """Test that BroadcastMarker is handled in paths (covers graph.py:698)."""
+    """Test that BroadcastMarker is handled in paths"""
     g = GraphBuilder(state_type=ExecutionState, output_type=list[str])
 
     @g.step
@@ -104,7 +104,7 @@ async def branch_b(ctx: StepContext[ExecutionState, None, str]) -> str:
 
 
 async def test_nested_joins_with_different_fork_stacks():
-    """Test nested joins with different fork stack depths (covers graph.py:556)."""
+    """Test nested joins with different fork stack depths"""
     g = GraphBuilder(state_type=ExecutionState, output_type=list[int])
 
     @g.step
@@ -188,7 +188,7 @@ async def test_empty_map_handling():
 
 
 async def test_complex_fork_stack_with_multiple_levels():
-    """Test complex scenarios with multiple fork levels (covers various graph.py lines)."""
+    """Test complex scenarios with multiple fork levels"""
     g = GraphBuilder(state_type=ExecutionState, output_type=list[int])
 
     @g.step
@@ -259,6 +259,45 @@ async def path_c(ctx: StepContext[ExecutionState, None, int]) -> int:
     assert sorted(result) == [20, 30, 40]
 
 
+async def test_implicit_broadcast_with_immediate_join():
+    """Test broadcast that immediately joins by just manually adding multiple edges from a single node."""
+    g = GraphBuilder(state_type=ExecutionState, output_type=list[int])
+
+    @g.step
+    async def source(ctx: StepContext[ExecutionState, None, None]) -> int:
+        return 10
+
+    @g.step
+    async def path_a(ctx: StepContext[ExecutionState, None, int]) -> int:
+        return ctx.inputs * 2
+
+    @g.step
+    async def path_b(ctx: StepContext[ExecutionState, None, int]) -> int:
+        return ctx.inputs * 3
+
+    @g.step
+    async def path_c(ctx: StepContext[ExecutionState, None, int]) -> int:
+        return ctx.inputs * 4
+
+    collect = g.join(reduce_list_append, initial_factory=list[int])
+
+    g.add(
+        g.edge_from(g.start_node).to(source),
+        # Multiple .to() destinations creates a broadcast
+        g.edge_from(source).to(path_a),
+        g.edge_from(source).to(path_b),
+        g.edge_from(source).to(path_c),
+        g.edge_from(path_a).to(collect),
+        g.edge_from(path_b).to(collect),
+        g.edge_from(path_c).to(collect),
+        g.edge_from(collect).to(g.end_node),
+    )
+
+    graph = g.build()
+    result = await graph.run(state=ExecutionState())
+    assert sorted(result) == [20, 30, 40]
+
+
 async def test_mixed_sequential_and_parallel_execution():
     """Test graph with both sequential and parallel sections."""
     g = GraphBuilder(state_type=ExecutionState, output_type=str)

tests/graph/beta/test_joins_and_reducers.py

@@ -262,7 +262,7 @@ async def create_dict(ctx: StepContext[SimpleState, None, int]) -> dict[str, int
 
 
 async def test_reducer_with_deps_access():
-    """Test that reducer context can access deps (covers join.py:71)."""
+    """Test that reducer context can access deps"""
 
     @dataclass
     class DepsWithMultiplier:
@@ -300,7 +300,7 @@ async def process(ctx: StepContext[SimpleState, DepsWithMultiplier, int]) -> int
 
 
 async def test_reduce_list_extend():
-    """Test reduce_list_extend that extends a list with iterables (covers join.py:115)."""
+    """Test reduce_list_extend that extends a list with iterables"""
     from pydantic_graph.beta.join import reduce_list_extend
 
     g = GraphBuilder(state_type=SimpleState, output_type=list[int])