simplify

Author: alexander-alderman-webb

sentry_sdk/integrations/anthropic.py

@@ -135,6 +135,13 @@ def setup_once() -> None:
             MessageStreamManager.__enter__
         )
 
+        # Before https://github.com/anthropics/anthropic-sdk-python/commit/b1a1c0354a9aca450a7d512fdbdeb59c0ead688a
+        # MessageStream inherits from Stream, so patching Stream is sufficient on these versions.
+        if not issubclass(MessageStream, Stream):
+            MessageStream.__iter__ = _wrap_message_stream_iter(MessageStream.__iter__)
+            MessageStream.__next__ = _wrap_message_stream_next(MessageStream.__next__)
+            MessageStream.close = _wrap_message_stream_close(MessageStream.close)
+
         AsyncMessages.stream = _wrap_async_message_stream(AsyncMessages.stream)
         AsyncMessageStreamManager.__aenter__ = (
             _wrap_async_message_stream_manager_aenter(
@@ -143,12 +150,8 @@ def setup_once() -> None:
         )
 
         # Before https://github.com/anthropics/anthropic-sdk-python/commit/b1a1c0354a9aca450a7d512fdbdeb59c0ead688a
-        # MessageStream inherits from Stream, so patching Stream is sufficient on these versions.
-        if not issubclass(MessageStream, Stream):
-            MessageStream.__iter__ = _wrap_message_stream_iter(MessageStream.__iter__)
-            MessageStream.__next__ = _wrap_message_stream_next(MessageStream.__next__)
-            MessageStream.close = _wrap_message_stream_close(MessageStream.close)
-
+        # AsyncMessageStream inherits from AsyncStream, so patching Stream is sufficient on these versions.
+        if not issubclass(AsyncMessageStream, AsyncStream):
             AsyncMessageStream.__aiter__ = _wrap_async_message_stream_aiter(
                 AsyncMessageStream.__aiter__
             )

tests/integrations/anthropic/test_anthropic.py

@@ -2,7 +2,6 @@
 from unittest import mock
 import json
 from itertools import islice
-from builtins import anext
 
 try:
     from unittest.mock import AsyncMock
@@ -1049,7 +1048,7 @@ async def test_streaming_create_message_async_next_consumption(
             )
 
             while True:
-                await anext(messages)
+                await messages.__anext__()
 
     assert len(events) == 1
     (event,) = events
@@ -1151,13 +1150,13 @@ async def test_streaming_create_message_async_iterator_methods(
                 max_tokens=1024, messages=messages, model="model", stream=True
             )
 
-            await anext(messages)
-            await anext(messages)
+            await messages.__anext__()
+            await messages.__anext__()
 
             async for item in messages:
                 break
 
-            await anext(messages)
+            await messages.__anext__()
             await messages.close()
 
     assert len(events) == 1
@@ -1383,7 +1382,7 @@ async def test_stream_messages_async_next_consumption(
                 model="model",
             ) as stream:
                 while True:
-                    await anext(stream)
+                    await stream.__anext__()
 
     assert len(events) == 1
     (event,) = events
@@ -1486,16 +1485,18 @@ async def test_stream_messages_async_iterator_methods(
                 messages=messages,
                 model="model",
             ) as stream:
-                await anext(stream)
-                await anext(stream)
+                await stream.__anext__()
+                await stream.__anext__()
 
                 async for item in stream:
                     break
 
-                await anext(stream)
+                await stream.__anext__()
                 # New versions add TextEvent, so consume one more event.
-                if TextEvent is not None and isinstance(await anext(stream), TextEvent):
-                    await anext(stream)
+                if TextEvent is not None and isinstance(
+                    await stream.__anext__(), TextEvent
+                ):
+                    await stream.__anext__()
                 await stream.close()
 
     assert len(events) == 1