Skip to content

Commit 6812235

Browse files
improve types
1 parent fe7b26b commit 6812235

1 file changed

Lines changed: 13 additions & 8 deletions

File tree

‎sentry_sdk/integrations/anthropic.py‎

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)