@@ -162,6 +162,7 @@ def response_hook(span, request, response):
162162"""
163163import logging
164164import typing
165+ from types import TracebackType
165166
166167import httpx
167168
@@ -293,6 +294,18 @@ def __init__(
293294 self ._request_hook = request_hook
294295 self ._response_hook = response_hook
295296
297+ def __enter__ (self ) -> "SyncOpenTelemetryTransport" :
298+ self ._transport .__enter__ ()
299+ return self
300+
301+ def __exit__ (
302+ self ,
303+ exc_type : typing .Optional [typing .Type [BaseException ]] = None ,
304+ exc_value : typing .Optional [BaseException ] = None ,
305+ traceback : typing .Optional [TracebackType ] = None ,
306+ ) -> None :
307+ self ._transport .__exit__ (exc_type , exc_value , traceback )
308+
296309 def handle_request (
297310 self ,
298311 * args ,
@@ -343,6 +356,9 @@ def handle_request(
343356
344357 return response
345358
359+ def close (self ) -> None :
360+ self ._transport .close ()
361+
346362
347363class AsyncOpenTelemetryTransport (httpx .AsyncBaseTransport ):
348364 """Async transport class that will trace all requests made with a client.
@@ -372,6 +388,18 @@ def __init__(
372388 self ._request_hook = request_hook
373389 self ._response_hook = response_hook
374390
391+ async def __aenter__ (self ) -> "AsyncOpenTelemetryTransport" :
392+ await self ._transport .__aenter__ ()
393+ return self
394+
395+ async def __aexit__ (
396+ self ,
397+ exc_type : typing .Optional [typing .Type [BaseException ]] = None ,
398+ exc_value : typing .Optional [BaseException ] = None ,
399+ traceback : typing .Optional [TracebackType ] = None ,
400+ ) -> None :
401+ await self ._transport .__aexit__ (exc_type , exc_value , traceback )
402+
375403 async def handle_async_request (
376404 self , * args , ** kwargs
377405 ) -> typing .Union [
@@ -423,6 +451,9 @@ async def handle_async_request(
423451
424452 return response
425453
454+ async def aclose (self ) -> None :
455+ await self ._transport .aclose ()
456+
426457
427458class _InstrumentedClient (httpx .Client ):
428459 _tracer_provider = None
0 commit comments