Add stream methods to httpx OAuth2 Clients
This implements OAuth2 for streaming requests, which aren't handled
by the request() methods of httpx, but have a separate stream()
method.
@@ -94,6 +94,19 @@ class AsyncOAuth2Client(_OAuth2Client, AsyncClient):
return await super(AsyncOAuth2Client, self).request(
method, url, auth=auth, **kwargs)
+ async def stream(self, method, url, withhold_token=False, auth=UNSET, **kwargs):
+ if not withhold_token and auth is UNSET:
+ if not self.token:
+ raise MissingTokenError()
+
+ if self.token.is_expired():
+ await self.ensure_active_token(self.token)
+
+ auth = self.token_auth
+
+ return super(AsyncOAuth2Client, self).stream(
+ method, url, auth=auth, **kwargs)
+
async def ensure_active_token(self, token):
if self._token_refresh_event.is_set():
# Unset the event so other coroutines don't try to update the token
@@ -199,3 +212,16 @@ class OAuth2Client(_OAuth2Client, Client):
return super(OAuth2Client, self).request(
method, url, auth=auth, **kwargs)
+
+ def stream(self, method, url, withhold_token=False, auth=UNSET, **kwargs):
+ if not withhold_token and auth is UNSET:
+ if not self.token:
+ raise MissingTokenError()
+
+ if not self.ensure_active_token(self.token):
+ raise InvalidTokenError()
+
+ auth = self.token_auth
+
+ return super(OAuth2Client, self).stream(
+ method, url, auth=auth, **kwargs)
