Refactor to use anyio

Author: dmontagu

pydantic_ai_slim/pydantic_ai/_utils.py

@@ -147,7 +147,7 @@ async def group_by_temporal(
         aiterable: The async iterable to group.
         soft_max_interval: Maximum interval over which to group items, this should avoid a trickle of items causing
             a group to never be yielded. It's a soft max in the sense that once we're over this time, we yield items
-            as soon as `aiter.__anext__()` returns. If `None`, no grouping/debouncing is performed
+            as soon as `anext(aiter)` returns. If `None`, no grouping/debouncing is performed
 
     Returns:
         A context manager usable as an async iterable of lists of items produced by the input async iterable.
@@ -171,7 +171,7 @@ async def async_iter_groups() -> AsyncIterator[list[T]]:
         buffer: list[T] = []
         group_start_time = time.monotonic()
 
-        aiterator = aiterable.__aiter__()
+        aiterator = aiter(aiterable)
         while True:
             if group_start_time is None:
                 # group hasn't started, we just wait for the maximum interval
@@ -182,9 +182,9 @@ async def async_iter_groups() -> AsyncIterator[list[T]]:
 
             # if there's no current task, we get the next one
             if task is None:
-                # aiter.__anext__() returns an Awaitable[T], not a Coroutine which asyncio.create_task expects
+                # anext(aiter) returns an Awaitable[T], not a Coroutine which asyncio.create_task expects
                 # so far, this doesn't seem to be a problem
-                task = asyncio.create_task(aiterator.__anext__())  # pyright: ignore[reportArgumentType]
+                task = asyncio.create_task(anext(aiterator))  # pyright: ignore[reportArgumentType]
 
             # we use asyncio.wait to avoid cancelling the coroutine if it's not done
             done, _ = await asyncio.wait((task,), timeout=wait_time)
@@ -284,10 +284,10 @@ async def peek(self) -> T | Unset:
 
         # Otherwise, we need to fetch the next item from the underlying iterator.
         if self._source_iter is None:
-            self._source_iter = self._source.__aiter__()
+            self._source_iter = aiter(self._source)
 
         try:
-            self._buffer = await self._source_iter.__anext__()
+            self._buffer = await anext(self._source_iter)
         except StopAsyncIteration:
             self._exhausted = True
             return UNSET
@@ -318,10 +318,10 @@ async def __anext__(self) -> T:
 
         # Otherwise, fetch the next item from the source.
         if self._source_iter is None:
-            self._source_iter = self._source.__aiter__()
+            self._source_iter = aiter(self._source)
 
         try:
-            return await self._source_iter.__anext__()
+            return await anext(self._source_iter)
         except StopAsyncIteration:
             self._exhausted = True
             raise

pydantic_ai_slim/pydantic_ai/run.py

@@ -145,7 +145,7 @@ async def __anext__(
         self,
     ) -> _agent_graph.AgentNode[AgentDepsT, OutputDataT] | End[FinalResult[OutputDataT]]:
         """Advance to the next node automatically based on the last returned node."""
-        task = await self._graph_run.__anext__()
+        task = await anext(self._graph_run)
         return self._task_to_node(task)
 
     def _task_to_node(

pydantic_graph/pydantic_graph/beta/decision.py

@@ -8,7 +8,7 @@
 from __future__ import annotations
 
 import inspect
-from collections.abc import Callable, Iterable, Sequence
+from collections.abc import AsyncIterable, Callable, Iterable, Sequence
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Generic, get_origin
 
@@ -232,7 +232,8 @@ def transform(
         )
 
     def map(
-        self: DecisionBranchBuilder[StateT, DepsT, Iterable[T], SourceT, HandledT],
+        self: DecisionBranchBuilder[StateT, DepsT, Iterable[T], SourceT, HandledT]
+        | DecisionBranchBuilder[StateT, DepsT, AsyncIterable[T], SourceT, HandledT],
         *,
         fork_id: str | None = None,
         downstream_join_id: str | None = None,