2020from __future__ import annotations
2121import logging
2222from typing import Any , Dict , Generator , Optional , Tuple , Union
23- import httpx
23+ import httpx2
2424import requests .adapters
2525from requests .packages .urllib3 .util import retry # pylint: disable=import-error
2626from google .auth import credentials
@@ -165,7 +165,7 @@ def __init__(self, **kwargs):
165165 def parse_body (self , resp ):
166166 return resp .json ()
167167
168- class GoogleAuthCredentialFlow (httpx .Auth ):
168+ class GoogleAuthCredentialFlow (httpx2 .Auth ):
169169 """Google Auth Credential Auth Flow"""
170170 def __init__ (self , credential : credentials .Credentials ):
171171 self ._credential = credential
@@ -174,7 +174,7 @@ def __init__(self, credential: credentials.Credentials):
174174
175175 def apply_auth_headers (
176176 self ,
177- request : httpx .Request ,
177+ request : httpx2 .Request ,
178178 auth_request : google_auth_requests .Request
179179 ) -> None :
180180 """A helper function that refreshes credentials if needed and mutates the request headers
@@ -189,7 +189,9 @@ def apply_auth_headers(
189189 )
190190 logger .debug ('Auth headers applied. Credential validity after: %s' , self ._credential .valid )
191191
192- def auth_flow (self , request : httpx .Request ) -> Generator [httpx .Request , httpx .Response , None ]:
192+ def auth_flow (
193+ self , request : httpx2 .Request
194+ ) -> Generator [httpx2 .Request , httpx2 .Response , None ]:
193195 _original_headers = request .headers .copy ()
194196 _credential_refresh_attempt = 0
195197
@@ -207,7 +209,7 @@ def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Re
207209 'Dispatching request, attempt %d of %d' ,
208210 _credential_refresh_attempt , self ._max_refresh_attempts
209211 )
210- response : httpx .Response = yield request
212+ response : httpx2 .Response = yield request
211213
212214 if response .status_code in self ._refresh_status_codes :
213215 if _credential_refresh_attempt < self ._max_refresh_attempts :
@@ -235,7 +237,7 @@ def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Re
235237 response .status_code
236238 )
237239 break
238- # The last yielded response is automatically returned by httpx 's auth flow.
240+ # The last yielded response is automatically returned by httpx2 's auth flow.
239241
240242class HttpxAsyncClient ():
241243 """Async HTTP client used to make HTTP/2 calls using HTTPX.
@@ -247,7 +249,7 @@ def __init__(
247249 self ,
248250 credential : Optional [credentials .Credentials ] = None ,
249251 base_url : str = '' ,
250- headers : Optional [Union [httpx .Headers , Dict [str , str ]]] = None ,
252+ headers : Optional [Union [httpx2 .Headers , Dict [str , str ]]] = None ,
251253 retry_config : HttpxRetry = DEFAULT_HTTPX_RETRY_CONFIG ,
252254 timeout : int = DEFAULT_TIMEOUT_SECONDS ,
253255 http2 : bool = True
@@ -279,15 +281,15 @@ def __init__(
279281 }
280282
281283 if credential :
282- self ._async_client = httpx .AsyncClient (
284+ self ._async_client = httpx2 .AsyncClient (
283285 http2 = http2 ,
284286 timeout = self ._timeout ,
285287 headers = self ._headers ,
286288 auth = GoogleAuthCredentialFlow (credential ), # Add auth flow for credentials.
287289 mounts = self ._mounts
288290 )
289291 else :
290- self ._async_client = httpx .AsyncClient (
292+ self ._async_client = httpx2 .AsyncClient (
291293 http2 = http2 ,
292294 timeout = self ._timeout ,
293295 headers = self ._headers ,
@@ -306,12 +308,12 @@ def timeout(self):
306308 def async_client (self ):
307309 return self ._async_client
308310
309- async def request (self , method : str , url : str , ** kwargs : Any ) -> httpx .Response :
311+ async def request (self , method : str , url : str , ** kwargs : Any ) -> httpx2 .Response :
310312 """Makes an HTTP call using the HTTPX library.
311313
312314 This is the sole entry point to the HTTPX library. All other helper methods in this
313315 class call this method to send HTTP requests out. Refer to
314- https://www.python-httpx.org /api/ for more information on supported options
316+ https://pydantic.dev/docs/httpx2/api /api/ for more information on supported options
315317 and features.
316318
317319 Args:
@@ -332,12 +334,12 @@ class call this method to send HTTP requests out. Refer to
332334 resp = await self ._async_client .request (method , self .base_url + url , ** kwargs )
333335 return resp .raise_for_status ()
334336
335- async def headers (self , method : str , url : str , ** kwargs : Any ) -> httpx .Headers :
337+ async def headers (self , method : str , url : str , ** kwargs : Any ) -> httpx2 .Headers :
336338 resp = await self .request (method , url , ** kwargs )
337339 return resp .headers
338340
339341 async def body_and_response (
340- self , method : str , url : str , ** kwargs : Any ) -> Tuple [Any , httpx .Response ]:
342+ self , method : str , url : str , ** kwargs : Any ) -> Tuple [Any , httpx2 .Response ]:
341343 resp = await self .request (method , url , ** kwargs )
342344 return self .parse_body (resp ), resp
343345
@@ -346,11 +348,11 @@ async def body(self, method: str, url: str, **kwargs: Any) -> Any:
346348 return self .parse_body (resp )
347349
348350 async def headers_and_body (
349- self , method : str , url : str , ** kwargs : Any ) -> Tuple [httpx .Headers , Any ]:
351+ self , method : str , url : str , ** kwargs : Any ) -> Tuple [httpx2 .Headers , Any ]:
350352 resp = await self .request (method , url , ** kwargs )
351353 return resp .headers , self .parse_body (resp )
352354
353- def parse_body (self , resp : httpx .Response ) -> Any :
355+ def parse_body (self , resp : httpx2 .Response ) -> Any :
354356 return resp .json ()
355357
356358 async def aclose (self ) -> None :
0 commit comments