From ca5d9a6bd7ed1b0610c68d3702985bcb02310323 Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Wed, 15 Apr 2026 19:21:08 -0400 Subject: [PATCH] Fix mypy errors, tighten init() return type, move import to top level --- libs/langgraph/langgraph/stream/_mux.py | 4 +--- libs/langgraph/langgraph/stream/_types.py | 2 +- libs/langgraph/langgraph/stream/run_stream.py | 15 ++++++--------- libs/langgraph/langgraph/stream/stream_channel.py | 3 +-- libs/langgraph/tests/test_streaming_handler.py | 4 +++- 5 files changed, 12 insertions(+), 16 deletions(-) diff --git a/libs/langgraph/langgraph/stream/_mux.py b/libs/langgraph/langgraph/stream/_mux.py index 3865f7c19..71a987fcf 100644 --- a/libs/langgraph/langgraph/stream/_mux.py +++ b/libs/langgraph/langgraph/stream/_mux.py @@ -117,9 +117,7 @@ class StreamMux: # Ensure the channel's log matches the mux's mode. if value._is_async != self._is_async: value._is_async = self._is_async - value._log = ( - AsyncEventLog() if self._is_async else EventLog() - ) + value._log = AsyncEventLog() if self._is_async else EventLog() self._channels.append(value) channel_name = value.name diff --git a/libs/langgraph/langgraph/stream/_types.py b/libs/langgraph/langgraph/stream/_types.py index 560acd046..3aec9c12f 100644 --- a/libs/langgraph/langgraph/stream/_types.py +++ b/libs/langgraph/langgraph/stream/_types.py @@ -45,7 +45,7 @@ class StreamTransformer(ABC): """ @abstractmethod - def init(self) -> Any: + def init(self) -> dict[str, Any]: """Return the projection dict. Keys become entries in `run.extensions`. If the transformer has diff --git a/libs/langgraph/langgraph/stream/run_stream.py b/libs/langgraph/langgraph/stream/run_stream.py index 53bb02e7f..4aab11a99 100644 --- a/libs/langgraph/langgraph/stream/run_stream.py +++ b/libs/langgraph/langgraph/stream/run_stream.py @@ -4,7 +4,8 @@ import asyncio from collections.abc import AsyncIterator, Iterator from typing import Any -from langgraph.stream._event_log import EventLog +from langgraph.stream._convert import convert_to_protocol_event +from langgraph.stream._event_log import AsyncEventLog, EventLog from langgraph.stream._mux import StreamMux from langgraph.stream._types import ProtocolEvent from langgraph.stream.stream_channel import StreamChannel @@ -43,18 +44,14 @@ class GraphRunStream: # when its cursor catches up to the buffer. self._wire_request_more(mux, extensions) - def _wire_request_more( - self, mux: StreamMux, extensions: dict[str, Any] - ) -> None: + def _wire_request_more(self, mux: StreamMux, extensions: dict[str, Any]) -> None: """Set _request_more on all sync EventLogs so iteration drives the graph.""" if isinstance(mux._events, EventLog): mux._events._request_more = self._pump_next for value in extensions.values(): if isinstance(value, EventLog): value._request_more = self._pump_next - elif isinstance(value, StreamChannel) and isinstance( - value._log, EventLog - ): + elif isinstance(value, StreamChannel) and isinstance(value._log, EventLog): value._log._request_more = self._pump_next def _pump_next(self) -> bool: @@ -74,8 +71,6 @@ class GraphRunStream: self._mux.fail(e) self._exhausted = True return False - from langgraph.stream._convert import convert_to_protocol_event - self._mux.push(convert_to_protocol_event(part)) return True @@ -106,6 +101,7 @@ class GraphRunStream: def __iter__(self) -> Iterator[ProtocolEvent]: """Iterate all protocol events from the mux's main event log.""" + assert isinstance(self._mux._events, EventLog) return iter(self._mux._events) @@ -172,4 +168,5 @@ class AsyncGraphRunStream: def __aiter__(self) -> AsyncIterator[ProtocolEvent]: """Iterate all protocol events from the mux's main event log.""" + assert isinstance(self._mux._events, AsyncEventLog) return self._mux._events.__aiter__() diff --git a/libs/langgraph/langgraph/stream/stream_channel.py b/libs/langgraph/langgraph/stream/stream_channel.py index 289297da8..45e552bd3 100644 --- a/libs/langgraph/langgraph/stream/stream_channel.py +++ b/libs/langgraph/langgraph/stream/stream_channel.py @@ -68,7 +68,6 @@ class StreamChannel(Generic[T]): def __aiter__(self) -> AsyncIterator[T]: if not isinstance(self._log, AsyncEventLog): raise RuntimeError( - "Cannot use async iteration on a sync StreamChannel. " - "Use 'for' instead." + "Cannot use async iteration on a sync StreamChannel. Use 'for' instead." ) return self._log.__aiter__() diff --git a/libs/langgraph/tests/test_streaming_handler.py b/libs/langgraph/tests/test_streaming_handler.py index a46a99e54..71e08df04 100644 --- a/libs/langgraph/tests/test_streaming_handler.py +++ b/libs/langgraph/tests/test_streaming_handler.py @@ -418,7 +418,6 @@ class TestStreamingHandlerSyncErrors: list(run) - class TestStreamingHandlerSyncInterrupt: def test_interrupted(self) -> None: graph = _build_interrupt_graph() @@ -773,6 +772,7 @@ class TestStreamMuxResilience: def test_close_continues_after_finalize_error(self) -> None: """If a transformer's finalize() raises, the main event log and remaining transformers should still be closed/finalized.""" + class BrokenFinalizer(StreamTransformer): def init(self) -> dict[str, Any]: return {} @@ -812,6 +812,7 @@ class TestStreamMuxResilience: def test_fail_continues_after_transformer_error(self) -> None: """If a transformer's fail() raises, the main event log and remaining transformers should still be failed.""" + class BrokenFailer(StreamTransformer): def init(self) -> dict[str, Any]: return {} @@ -848,6 +849,7 @@ class TestStreamMuxResilience: def test_close_still_closes_channels_after_finalize_error(self) -> None: """Channels should be closed even if a transformer's finalize raises.""" + class BrokenWithChannel(StreamTransformer): def __init__(self) -> None: self._channel: StreamChannel[str] = StreamChannel("ch")