Try to improve coverage

Author: dmontagu

pydantic_graph/pydantic_graph/beta/graph.py

@@ -84,10 +84,6 @@ def __init__(self, value: OutputT):
     def value(self) -> OutputT:
         return self._value
 
-    @property
-    def end_marker_id(self) -> str:
-        return f'EndMarker:{id(self)}'
-
 
 @dataclass
 class JoinItem:
@@ -493,8 +489,6 @@ class _GraphIterator(Generic[StateT, DepsT, OutputT]):
 
     cancel_scopes: dict[TaskID, CancelScope] = field(init=False)
     active_tasks: dict[TaskID, GraphTask] = field(init=False)
-    pending_task_results: set[TaskID] = field(init=False)
-    cancelled_tasks: set[TaskID] = field(init=False)
     active_reducers: dict[tuple[JoinID, NodeRunID], JoinState] = field(init=False)
     iter_stream_sender: MemoryObjectSendStream[_GraphTaskResult] = field(init=False)
     iter_stream_receiver: MemoryObjectReceiveStream[_GraphTaskResult] = field(init=False)
@@ -506,9 +500,6 @@ def __post_init__(self):
         self.active_reducers = {}
         self.iter_stream_sender, self.iter_stream_receiver = create_memory_object_stream[_GraphTaskResult]()
 
-        self.pending_task_results = set()
-        self.cancelled_tasks = set()
-
     @property
     def task_group(self) -> TaskGroup:
         if self._task_group is None:
@@ -529,21 +520,6 @@ async def iter_graph(  # noqa C901
                     async with self.iter_stream_receiver:
                         while self.active_tasks or self.active_reducers:
                             async for task_result in self.iter_stream_receiver:  # pragma: no branch
-                                # If we encounter a mock task, add it to the active tasks to ensure we don't proceed until everything downstream is handled
-                                if (
-                                    not task_result.source_is_finished
-                                    and task_result.source.task_id not in self.active_tasks
-                                ):
-                                    self.active_tasks[task_result.source.task_id] = task_result.source
-
-                                if task_result.source.task_id in self.cancelled_tasks:
-                                    if task_result.source_is_finished:
-                                        self.cancelled_tasks.remove(task_result.source.task_id)
-                                    continue
-
-                                if task_result.source_is_finished:
-                                    self.pending_task_results.discard(task_result.source.task_id)
-
                                 if isinstance(task_result.result, JoinItem):
                                     maybe_overridden_result = task_result.result
                                 else:
@@ -574,7 +550,7 @@ async def iter_graph(  # noqa C901
                                     join_state.current = join_node.reduce(context, join_state.current, result.inputs)
                                     if join_state.cancelled_sibling_tasks:
                                         await self._cancel_sibling_tasks(parent_fork_id, fork_run_id)
-                                    if task_result.source_is_finished:
+                                    if task_result.source_is_finished:  # pragma: no branch
                                         await self._finish_task(task_result.source.task_id)
                                 else:
                                     for new_task in maybe_overridden_result:
@@ -636,13 +612,18 @@ async def iter_graph(  # noqa C901
                                         join_node, join_state.current, join_state.downstream_fork_stack
                                     )
                                     maybe_overridden_result = yield new_tasks
-                                    if isinstance(maybe_overridden_result, EndMarker):
+                                    if isinstance(maybe_overridden_result, EndMarker):  # pragma: no cover
+                                        # This is theoretically reachable but it would be awkward.
+                                        # Probably a better way to get coverage here would be to unify the code pat
+                                        # with the other `if isinstance(maybe_overridden_result, EndMarker):`
                                         self.task_group.cancel_scope.cancel()
                                         return
                                     for new_task in maybe_overridden_result:
                                         self.active_tasks[new_task.task_id] = new_task
                                     new_task_ids = {t.task_id for t in maybe_overridden_result}
-                                    for t in new_tasks:
+                                    for t in new_tasks:  # pragma: no cover
+                                        # Same note as above about how this is theoretically reachable but we should
+                                        # just get coverage by unifying the code paths
                                         if t.task_id not in new_task_ids:
                                             await self._finish_task(t.task_id)
                                     self._handle_execution_request(maybe_overridden_result)
@@ -674,7 +655,6 @@ async def _run_tracked_task(self, t_: GraphTask):
                     await self.iter_stream_sender.send(_GraphTaskResult(t_, new_tasks, False))
                 await self.iter_stream_sender.send(_GraphTaskResult(t_, []))
             else:
-                self.pending_task_results.add(t_.task_id)
                 await self.iter_stream_sender.send(_GraphTaskResult(t_, result))
 
     async def _run_task(
@@ -872,8 +852,6 @@ async def _cancel_sibling_tasks(self, parent_fork_id: ForkID, node_run_id: NodeR
                 else:
                     pass
         for task_id in task_ids_to_cancel:
-            if task_id in self.pending_task_results:
-                self.cancelled_tasks.add(task_id)
             await self._finish_task(task_id)
 
 

pydantic_graph/pydantic_graph/beta/graph_builder.py

@@ -44,7 +44,7 @@
     Path,
     PathBuilder,
 )
-from pydantic_graph.beta.step import NodeStep, Step, StepAsyncIteratorFunction, StepContext, StepFunction, StepNode
+from pydantic_graph.beta.step import NodeStep, Step, StepContext, StepFunction, StepNode, StreamFunction
 from pydantic_graph.beta.util import TypeOrTypeExpression, get_callable_name, unpack_type_expression
 from pydantic_graph.exceptions import GraphBuildingError
 from pydantic_graph.nodes import BaseNode, End
@@ -215,50 +215,50 @@ def decorator(
         return step
 
     @overload
-    def step_async_iterable(
+    def stream(
         self,
         *,
         node_id: str | None = None,
         label: str | None = None,
     ) -> Callable[
-        [StepAsyncIteratorFunction[StateT, DepsT, InputT, OutputT]], Step[StateT, DepsT, InputT, AsyncIterable[OutputT]]
+        [StreamFunction[StateT, DepsT, InputT, OutputT]], Step[StateT, DepsT, InputT, AsyncIterable[OutputT]]
     ]: ...
     @overload
-    def step_async_iterable(
+    def stream(
         self,
-        call: StepAsyncIteratorFunction[StateT, DepsT, InputT, OutputT],
+        call: StreamFunction[StateT, DepsT, InputT, OutputT],
         *,
         node_id: str | None = None,
         label: str | None = None,
     ) -> Step[StateT, DepsT, InputT, AsyncIterable[OutputT]]: ...
     @overload
-    def step_async_iterable(
+    def stream(
         self,
-        call: StepAsyncIteratorFunction[StateT, DepsT, InputT, OutputT] | None = None,
+        call: StreamFunction[StateT, DepsT, InputT, OutputT] | None = None,
         *,
         node_id: str | None = None,
         label: str | None = None,
     ) -> (
         Step[StateT, DepsT, InputT, AsyncIterable[OutputT]]
         | Callable[
-            [StepAsyncIteratorFunction[StateT, DepsT, InputT, OutputT]],
+            [StreamFunction[StateT, DepsT, InputT, OutputT]],
             Step[StateT, DepsT, InputT, AsyncIterable[OutputT]],
         ]
     ): ...
-    def step_async_iterable(
+    def stream(
         self,
-        call: StepAsyncIteratorFunction[StateT, DepsT, InputT, OutputT] | None = None,
+        call: StreamFunction[StateT, DepsT, InputT, OutputT] | None = None,
         *,
         node_id: str | None = None,
         label: str | None = None,
     ) -> (
         Step[StateT, DepsT, InputT, AsyncIterable[OutputT]]
         | Callable[
-            [StepAsyncIteratorFunction[StateT, DepsT, InputT, OutputT]],
+            [StreamFunction[StateT, DepsT, InputT, OutputT]],
             Step[StateT, DepsT, InputT, AsyncIterable[OutputT]],
         ]
     ):
-        """Create a step from a step function.
+        """Create a step from an async iterator (which functions like a "stream").
 
         This method can be used as a decorator or called directly to create
         a step node from an async function.
@@ -274,9 +274,9 @@ def step_async_iterable(
         if call is None:
 
             def decorator(
-                func: StepAsyncIteratorFunction[StateT, DepsT, InputT, OutputT],
+                func: StreamFunction[StateT, DepsT, InputT, OutputT],
             ) -> Step[StateT, DepsT, InputT, AsyncIterable[OutputT]]:
-                return self.step_async_iterable(call=func, node_id=node_id, label=label)
+                return self.stream(call=func, node_id=node_id, label=label)
 
             return decorator
 

pydantic_graph/pydantic_graph/beta/join.py

@@ -70,10 +70,6 @@ def state(self) -> StateT:
     def deps(self) -> DepsT:
         return self._deps
 
-    @property
-    def cancelled_sibling_tasks(self):
-        return self._join_state.cancelled_sibling_tasks
-
     def cancel_sibling_tasks(self):
         self._join_state.cancelled_sibling_tasks = True
 

pydantic_graph/pydantic_graph/beta/step.py

@@ -67,9 +67,7 @@ def inputs(self) -> InputT:
 class StepFunction(Protocol[StateT, DepsT, InputT, OutputT]):
     """Protocol for step functions that can be executed in the graph.
 
-    Step functions are async callables that receive a step context and return
-    a result. This protocol enables serialization and deserialization of step
-    calls similar to how evaluators work.
+    Step functions are async callables that receive a step context and return a result.
 
     Type Parameters:
         StateT: The type of the graph state
@@ -90,12 +88,10 @@ def __call__(self, ctx: StepContext[StateT, DepsT, InputT]) -> Awaitable[OutputT
         raise NotImplementedError
 
 
-class StepAsyncIteratorFunction(Protocol[StateT, DepsT, InputT, OutputT]):
-    """Protocol for step functions that can be executed in the graph.
+class StreamFunction(Protocol[StateT, DepsT, InputT, OutputT]):
+    """Protocol for stream functions that can be executed in the graph.
 
-    Step functions are async callables that receive a step context and return
-    a result. This protocol enables serialization and deserialization of step
-    calls similar to how evaluators work.
+    Stream functions are async callables that receive a step context and return an async iterator.
 
     Type Parameters:
         StateT: The type of the graph state
@@ -105,13 +101,13 @@ class StepAsyncIteratorFunction(Protocol[StateT, DepsT, InputT, OutputT]):
     """
 
     def __call__(self, ctx: StepContext[StateT, DepsT, InputT]) -> AsyncIterator[OutputT]:
-        """Execute the step function with the given context.
+        """Execute the stream function with the given context.
 
         Args:
             ctx: The step context containing state, dependencies, and inputs
 
         Returns:
-            An awaitable that resolves to the step's output
+            An async iterator yielding the streamed output
         """
         raise NotImplementedError
         yield

tests/graph/beta/test_graph_iteration.py

@@ -286,7 +286,7 @@ async def test_iter_with_async_iterable_map():
 
     g = GraphBuilder(state_type=IterState, output_type=list[int])
 
-    @g.step_async_iterable
+    @g.stream()
     async def generate_async(ctx: StepContext[IterState, None, None]) -> AsyncIterator[int]:
         for i in [1, 2, 3, 4]:
             yield i
@@ -346,21 +346,18 @@ async def process(ctx: StepContext[IterState, None, int]) -> int:
 
     async with graph.iter(state=state) as run:
         while True:
-            try:
-                event = await run.next()
-                if isinstance(event, list):
-                    # Filter out tasks where the node is 'process' and input is > 3
-                    filtered_tasks = [
-                        task
-                        for task in event
-                        if not (task.node_id == NodeID('process') and isinstance(task.inputs, int) and task.inputs > 3)
-                    ]
-                    if filtered_tasks != event:
-                        # Override with filtered tasks
-                        event = await run.next(filtered_tasks)
-                if isinstance(event, EndMarker):
-                    break
-            except StopAsyncIteration:
+            event = await run.next()
+            if isinstance(event, list):
+                # Filter out tasks where the node is 'process' and input is > 3
+                filtered_tasks = [
+                    task
+                    for task in event
+                    if not (task.node_id == NodeID('process') and isinstance(task.inputs, int) and task.inputs > 3)
+                ]
+                if filtered_tasks != event:
+                    # Override with filtered tasks
+                    event = await run.next(filtered_tasks)
+            if isinstance(event, EndMarker):
                 break
 
     # Only items <= 3 should have been processed
@@ -397,24 +394,21 @@ async def second_step(ctx: StepContext[IterState, None, int]) -> int:
     override_done = False
     async with graph.iter(state=state) as run:
         while True:
-            try:
-                event = await run.next()
-                if isinstance(event, EndMarker) and not override_done:
-                    # Instead of ending, create a new task
-                    # Get the fork_stack from the EndMarker's source
-                    fork_stack = run.next_task[0].fork_stack if isinstance(run.next_task, list) else ()
-
-                    new_task = GraphTask(
-                        node_id=NodeID('second_step'),
-                        inputs=event.value,
-                        fork_stack=fork_stack,
-                    )
-
-                    override_done = True
-                    event = await run.next([new_task])
-                if isinstance(event, EndMarker) and override_done:
-                    break
-            except StopAsyncIteration:
+            event = await run.next()
+            if isinstance(event, EndMarker) and not override_done:
+                # Instead of ending, create a new task
+                # Get the fork_stack from the EndMarker's source
+                fork_stack = run.next_task[0].fork_stack if isinstance(run.next_task, list) else ()
+
+                new_task = GraphTask(
+                    node_id=NodeID('second_step'),
+                    inputs=event.value,
+                    fork_stack=fork_stack,
+                )
+
+                override_done = True
+                event = await run.next([new_task])
+            if isinstance(event, EndMarker) and override_done:
                 break
 
     result = run.output
@@ -461,9 +455,7 @@ async def step3(ctx: StepContext[IterState, None, int]) -> str:  # pragma: no co
                     if any(task.node_id == NodeID('step2') for task in event):
                         # Override with an EndMarker to terminate early
                         early_exit_done = True
-                        event = await run.next(EndMarker('early_exit'))
-                if isinstance(event, EndMarker):
-                    break
+                        await run.next(EndMarker('early_exit'))
             except StopAsyncIteration:
                 break