@@ -88,7 +88,7 @@ async def _():
88
88
89
89
async def _stream (
90
90
self , method : Literal ["GET" ], url : str , * args , ** kwargs
91
- ) -> AsyncGenerator [bytes , None ]:
91
+ ) -> AsyncGenerator [httpx . Response , None ]:
92
92
n_attempts = self .configuration .retries .total
93
93
assert isinstance (n_attempts , int )
94
94
@@ -98,17 +98,15 @@ async def _stream(
98
98
stop = tenacity .stop_after_attempt (n_attempts ),
99
99
retry = tenacity .retry_if_exception_type (httpx .HTTPStatusError ),
100
100
)
101
- async def _ () -> AsyncGenerator [bytes , None ]:
101
+ async def _ () -> AsyncGenerator [httpx . Response , None ]:
102
102
async with self ._client .stream (
103
103
method = method , url = url , * args , ** kwargs
104
104
) as response :
105
105
if response .status_code in self .configuration .retries .status_forcelist :
106
106
response .raise_for_status ()
107
- async for chunk in response .aiter_bytes ():
108
- yield chunk
107
+ yield response
109
108
110
- async for chunk in _ ():
111
- yield chunk
109
+ return _ ()
112
110
113
111
async def put (self , * args , ** kwargs ) -> httpx .Response :
114
112
return await self ._request (self ._client .put , * args , ** kwargs )
@@ -127,9 +125,8 @@ async def get(self, *args, **kwargs) -> httpx.Response:
127
125
128
126
async def stream (
129
127
self , method : Literal ["GET" ], url : str , * args , ** kwargs
130
- ) -> AsyncGenerator [bytes , None ]:
131
- async for chunk in self ._stream (method = method , url = url , * args , ** kwargs ):
132
- yield chunk
128
+ ) -> AsyncGenerator [httpx .Response , None ]:
129
+ return await self ._stream (method = method , url = url , * args , ** kwargs )
133
130
134
131
def _wait_callback (self , retry_state : tenacity .RetryCallState ) -> int :
135
132
assert retry_state .outcome is not None
0 commit comments