1717not create spans on its own.
1818"""
1919
20+ from __future__ import annotations
21+
2022import contextlib
2123import http .client
2224import logging
2325import socket # pylint:disable=unused-import # Used for typing
2426import typing
25- from typing import Collection
27+ from typing import Any , Callable , Collection , TypedDict , cast
2628
2729import wrapt
2830
3638
3739logger = logging .getLogger (__name__ )
3840
41+ R = typing .TypeVar ("R" )
42+
3943
4044class HttpClientInstrumentor (BaseInstrumentor ):
4145 def instrumentation_dependencies (self ) -> Collection [str ]:
4246 return () # This instruments http.client from stdlib; no extra deps.
4347
44- def _instrument (self , ** kwargs ):
48+ def _instrument (self , ** kwargs : Any ):
4549 """Instruments the http.client module (not creating spans on its own)"""
4650 _instrument ()
4751
48- def _uninstrument (self , ** kwargs ):
52+ def _uninstrument (self , ** kwargs : Any ):
4953 _uninstrument ()
5054
5155
52- def _remove_nonrecording (spanlist : typing . List [Span ]):
56+ def _remove_nonrecording (spanlist : list [Span ]) -> bool :
5357 idx = len (spanlist ) - 1
5458 while idx >= 0 :
5559 if not spanlist [idx ].is_recording ():
@@ -67,7 +71,9 @@ def _remove_nonrecording(spanlist: typing.List[Span]):
6771 return True
6872
6973
70- def trysetip (conn : http .client .HTTPConnection , loglevel = logging .DEBUG ) -> bool :
74+ def trysetip (
75+ conn : http .client .HTTPConnection , loglevel : int = logging .DEBUG
76+ ) -> bool :
7177 """Tries to set the net.peer.ip semantic attribute on the current span from the given
7278 HttpConnection.
7379
@@ -110,14 +116,17 @@ def trysetip(conn: http.client.HTTPConnection, loglevel=logging.DEBUG) -> bool:
110116
111117
112118def _instrumented_connect (
113- wrapped , instance : http .client .HTTPConnection , args , kwargs
114- ):
119+ wrapped : Callable [..., R ],
120+ instance : http .client .HTTPConnection ,
121+ args : tuple [Any , ...],
122+ kwargs : dict [str , Any ],
123+ ) -> R :
115124 result = wrapped (* args , ** kwargs )
116125 trysetip (instance , loglevel = logging .WARNING )
117126 return result
118127
119128
120- def instrument_connect (module , name = "connect" ):
129+ def instrument_connect (module : type [ Any ] , name : str = "connect" ):
121130 """Instrument additional connect() methods, e.g. for derived classes."""
122131
123132 wrapt .wrap_function_wrapper (
@@ -129,8 +138,11 @@ def instrument_connect(module, name="connect"):
129138
130139def _instrument ():
131140 def instrumented_send (
132- wrapped , instance : http .client .HTTPConnection , args , kwargs
133- ):
141+ wrapped : Callable [..., R ],
142+ instance : http .client .HTTPConnection ,
143+ args : tuple [Any , ...],
144+ kwargs : dict [str , Any ],
145+ ) -> R :
134146 done = trysetip (instance )
135147 result = wrapped (* args , ** kwargs )
136148 if not done :
@@ -147,8 +159,12 @@ def instrumented_send(
147159 # No need to instrument HTTPSConnection, as it calls super().connect()
148160
149161
150- def _getstate () -> typing .Optional [dict ]:
151- return context .get_value (_STATE_KEY )
162+ class _ConnectionState (TypedDict ):
163+ need_ip : list [Span ]
164+
165+
166+ def _getstate () -> _ConnectionState | None :
167+ return cast (_ConnectionState , context .get_value (_STATE_KEY ))
152168
153169
154170@contextlib .contextmanager
@@ -163,7 +179,7 @@ def set_ip_on_next_http_connection(span: Span):
163179 finally :
164180 context .detach (token )
165181 else :
166- spans : typing . List [ Span ] = state ["need_ip" ]
182+ spans = state ["need_ip" ]
167183 spans .append (span )
168184 try :
169185 yield
0 commit comments