17
17
not create spans on its own.
18
18
"""
19
19
20
+ from __future__ import annotations
21
+
20
22
import contextlib
21
23
import http .client
22
24
import logging
23
25
import socket # pylint:disable=unused-import # Used for typing
24
26
import typing
25
- from typing import Collection
27
+ from typing import Any , Callable , Collection , TypedDict , cast
26
28
27
29
import wrapt
28
30
36
38
37
39
logger = logging .getLogger (__name__ )
38
40
41
+ R = typing .TypeVar ("R" )
42
+
39
43
40
44
class HttpClientInstrumentor (BaseInstrumentor ):
41
45
def instrumentation_dependencies (self ) -> Collection [str ]:
42
46
return () # This instruments http.client from stdlib; no extra deps.
43
47
44
- def _instrument (self , ** kwargs ):
48
+ def _instrument (self , ** kwargs : Any ):
45
49
"""Instruments the http.client module (not creating spans on its own)"""
46
50
_instrument ()
47
51
48
- def _uninstrument (self , ** kwargs ):
52
+ def _uninstrument (self , ** kwargs : Any ):
49
53
_uninstrument ()
50
54
51
55
52
- def _remove_nonrecording (spanlist : typing . List [Span ]):
56
+ def _remove_nonrecording (spanlist : list [Span ]) -> bool :
53
57
idx = len (spanlist ) - 1
54
58
while idx >= 0 :
55
59
if not spanlist [idx ].is_recording ():
@@ -67,7 +71,9 @@ def _remove_nonrecording(spanlist: typing.List[Span]):
67
71
return True
68
72
69
73
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 :
71
77
"""Tries to set the net.peer.ip semantic attribute on the current span from the given
72
78
HttpConnection.
73
79
@@ -110,14 +116,17 @@ def trysetip(conn: http.client.HTTPConnection, loglevel=logging.DEBUG) -> bool:
110
116
111
117
112
118
def _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 :
115
124
result = wrapped (* args , ** kwargs )
116
125
trysetip (instance , loglevel = logging .WARNING )
117
126
return result
118
127
119
128
120
- def instrument_connect (module , name = "connect" ):
129
+ def instrument_connect (module : type [ Any ] , name : str = "connect" ):
121
130
"""Instrument additional connect() methods, e.g. for derived classes."""
122
131
123
132
wrapt .wrap_function_wrapper (
@@ -129,8 +138,11 @@ def instrument_connect(module, name="connect"):
129
138
130
139
def _instrument ():
131
140
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 :
134
146
done = trysetip (instance )
135
147
result = wrapped (* args , ** kwargs )
136
148
if not done :
@@ -147,8 +159,12 @@ def instrumented_send(
147
159
# No need to instrument HTTPSConnection, as it calls super().connect()
148
160
149
161
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 ))
152
168
153
169
154
170
@contextlib .contextmanager
@@ -163,7 +179,7 @@ def set_ip_on_next_http_connection(span: Span):
163
179
finally :
164
180
context .detach (token )
165
181
else :
166
- spans : typing . List [ Span ] = state ["need_ip" ]
182
+ spans = state ["need_ip" ]
167
183
spans .append (span )
168
184
try :
169
185
yield
0 commit comments