@@ -285,11 +285,14 @@ def setup_once() -> None:
285285 if isinstance (Stream .__dict__ .get ("_iterator" ), cached_property ):
286286 _original_stream_iterator = Stream ._iterator .func # type: ignore[attr-defined]
287287
288- def _sentry_patched_stream_iterator (self : "Any " ) -> "Any" :
289- original = _original_stream_iterator (self )
288+ def _sentry_patched_stream_iterator (self : "Stream " ) -> "Any" :
289+ iterator = _original_stream_iterator (self )
290290 if not hasattr (self , "_span" ):
291- return original
292- return _wrap_synchronous_message_iterator (self , original )
291+ return iterator
292+
293+ return _wrap_synchronous_message_iterator (
294+ cast ("_PatchedRawMessageStream" , self ), iterator
295+ )
293296
294297 patched = cached_property (_sentry_patched_stream_iterator )
295298 patched .__set_name__ (Stream , "_iterator" )
@@ -298,11 +301,13 @@ def _sentry_patched_stream_iterator(self: "Any") -> "Any":
298301 if isinstance (AsyncStream .__dict__ .get ("_iterator" ), cached_property ):
299302 _original_async_stream_iterator = AsyncStream ._iterator .func # type: ignore[attr-defined]
300303
301- def _sentry_patched_async_stream_iterator (self : "Any " ) -> "Any" :
302- original = _original_async_stream_iterator (self )
304+ def _sentry_patched_async_stream_iterator (self : "AsyncStream " ) -> "Any" :
305+ iterator = _original_async_stream_iterator (self )
303306 if not hasattr (self , "_span" ):
304- return original
305- return _wrap_asynchronous_message_iterator (self , original )
307+ return iterator
308+ return _wrap_asynchronous_message_iterator (
309+ cast ("_PatchedRawAsyncMessageStream" , self ), iterator
310+ )
306311
307312 patched = cached_property (_sentry_patched_async_stream_iterator )
308313 patched .__set_name__ (AsyncStream , "_iterator" )
0 commit comments