Support PEP 561 to opentelemetry-instrumentation-wsgi

Author: Kludex

instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/__init__.py

@@ -97,15 +97,22 @@ def GET(self):
 
 .. code-block:: python
 
+    from wsgiref.types import WSGIEnvironment, StartResponse
+    from opentelemetry.instrumentation.wsgi import OpenTelemetryMiddleware
+
+    def app(environ: WSGIEnvironment, start_response: StartResponse):
+        start_response("200 OK", [("Content-Type", "text/plain"), ("Content-Length", "13")])
+        return [b"Hello, World!"]
+
     def request_hook(span: Span, environ: WSGIEnvironment):
         if span and span.is_recording():
             span.set_attribute("custom_user_attribute_from_request_hook", "some-value")
 
-    def response_hook(span: Span, environ: WSGIEnvironment, status: str, response_headers: List):
+    def response_hook(span: Span, environ: WSGIEnvironment, status: str, response_headers: list[tuple[str, str]]):
         if span and span.is_recording():
             span.set_attribute("custom_user_attribute_from_response_hook", "some-value")
 
-    OpenTelemetryMiddleware(request_hook=request_hook, response_hook=response_hook)
+    OpenTelemetryMiddleware(app, request_hook=request_hook, response_hook=response_hook)
 
 Capture HTTP request and response headers
 *****************************************
@@ -207,10 +214,13 @@ def response_hook(span: Span, environ: WSGIEnvironment, status: str, response_he
 ---
 """
 
+from __future__ import annotations
+
 import functools
-import typing
 import wsgiref.util as wsgiref_util
 from timeit import default_timer
+from typing import TYPE_CHECKING, Any, Callable, Iterable, TypeVar, cast
+from wsgiref.types import StartResponse
 
 from opentelemetry import context, trace
 from opentelemetry.instrumentation._semconv import (
@@ -240,14 +250,15 @@ def response_hook(span: Span, environ: WSGIEnvironment, status: str, response_he
 )
 from opentelemetry.instrumentation.utils import _start_internal_or_server_span
 from opentelemetry.instrumentation.wsgi.version import __version__
-from opentelemetry.metrics import get_meter
+from opentelemetry.metrics import MeterProvider, get_meter
 from opentelemetry.propagators.textmap import Getter
 from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
 from opentelemetry.semconv.metrics import MetricInstruments
 from opentelemetry.semconv.metrics.http_metrics import (
     HTTP_SERVER_REQUEST_DURATION,
 )
 from opentelemetry.semconv.trace import SpanAttributes
+from opentelemetry.trace import TracerProvider
 from opentelemetry.trace.status import Status, StatusCode
 from opentelemetry.util.http import (
     OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS,
@@ -262,15 +273,23 @@ def response_hook(span: Span, environ: WSGIEnvironment, status: str, response_he
     sanitize_method,
 )
 
+if TYPE_CHECKING:
+    from wsgiref.types import WSGIApplication, WSGIEnvironment
+
+
+T = TypeVar("T")
+RequestHook = Callable[[trace.Span, WSGIEnvironment], None]
+ResponseHook = Callable[
+    [trace.Span, WSGIEnvironment, str, list[tuple[str, str]]], None
+]
+
 _HTTP_VERSION_PREFIX = "HTTP/"
 _CARRIER_KEY_PREFIX = "HTTP_"
 _CARRIER_KEY_PREFIX_LEN = len(_CARRIER_KEY_PREFIX)
 
 
-class WSGIGetter(Getter[dict]):
-    def get(
-        self, carrier: dict, key: str
-    ) -> typing.Optional[typing.List[str]]:
+class WSGIGetter(Getter[dict[str, Any]]):
+    def get(self, carrier: dict[str, Any], key: str) -> list[str] | None:
         """Getter implementation to retrieve a HTTP header value from the
              PEP3333-conforming WSGI environ
 
@@ -287,7 +306,7 @@ def get(
             return [value]
         return None
 
-    def keys(self, carrier):
+    def keys(self, carrier: dict[str, Any]):
         return [
             key[_CARRIER_KEY_PREFIX_LEN:].lower().replace("_", "-")
             for key in carrier
@@ -298,26 +317,19 @@ def keys(self, carrier):
 wsgi_getter = WSGIGetter()
 
 
-def setifnotnone(dic, key, value):
-    if value is not None:
-        dic[key] = value
-
-
 # pylint: disable=too-many-branches
-
-
 def collect_request_attributes(
-    environ,
-    sem_conv_opt_in_mode=_StabilityMode.DEFAULT,
+    environ: WSGIEnvironment,
+    sem_conv_opt_in_mode: _StabilityMode = _StabilityMode.DEFAULT,
 ):
     """Collects HTTP request attributes from the PEP3333-conforming
     WSGI environ and returns a dictionary to be used as span creation attributes.
     """
-    result = {}
+    result: dict[str, str | None] = {}
     _set_http_method(
         result,
         environ.get("REQUEST_METHOD", ""),
-        sanitize_method(environ.get("REQUEST_METHOD", "")),
+        sanitize_method(cast(str, environ.get("REQUEST_METHOD", ""))),
         sem_conv_opt_in_mode,
     )
     # old semconv v1.12.0
@@ -385,7 +397,7 @@ def collect_request_attributes(
     return result
 
 
-def collect_custom_request_headers_attributes(environ):
+def collect_custom_request_headers_attributes(environ: WSGIEnvironment):
     """Returns custom HTTP request headers which are configured by the user
     from the PEP3333-conforming WSGI environ to be used as span creation attributes as described
     in the specification https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/http.md#http-request-and-response-headers
@@ -411,7 +423,9 @@ def collect_custom_request_headers_attributes(environ):
     )
 
 
-def collect_custom_response_headers_attributes(response_headers):
+def collect_custom_response_headers_attributes(
+    response_headers: list[tuple[str, str]],
+):
     """Returns custom HTTP response headers which are configured by the user from the
     PEP3333-conforming WSGI environ as described in the specification
     https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/http.md#http-request-and-response-headers
@@ -422,7 +436,7 @@ def collect_custom_response_headers_attributes(response_headers):
             OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS
         )
     )
-    response_headers_dict = {}
+    response_headers_dict: dict[str, str] = {}
     if response_headers:
         for key, val in response_headers:
             key = key.lower()
@@ -440,7 +454,8 @@ def collect_custom_response_headers_attributes(response_headers):
     )
 
 
-def _parse_status_code(resp_status):
+# TODO: Used only on the `opentelemetry-instrumentation-pyramid` package - It can be moved there.
+def _parse_status_code(resp_status: str) -> int | None:
     status_code, _ = resp_status.split(" ", 1)
     try:
         return int(status_code)
@@ -449,7 +464,7 @@ def _parse_status_code(resp_status):
 
 
 def _parse_active_request_count_attrs(
-    req_attrs, sem_conv_opt_in_mode=_StabilityMode.DEFAULT
+    req_attrs, sem_conv_opt_in_mode: _StabilityMode = _StabilityMode.DEFAULT
 ):
     return _filter_semconv_active_request_count_attr(
         req_attrs,
@@ -460,7 +475,8 @@ def _parse_active_request_count_attrs(
 
 
 def _parse_duration_attrs(
-    req_attrs, sem_conv_opt_in_mode=_StabilityMode.DEFAULT
+    req_attrs: dict[str, str | None],
+    sem_conv_opt_in_mode: _StabilityMode = _StabilityMode.DEFAULT,
 ):
     return _filter_semconv_duration_attrs(
         req_attrs,
@@ -471,11 +487,11 @@ def _parse_duration_attrs(
 
 
 def add_response_attributes(
-    span,
-    start_response_status,
-    response_headers,
-    duration_attrs=None,
-    sem_conv_opt_in_mode=_StabilityMode.DEFAULT,
+    span: trace.Span,
+    start_response_status: str,
+    response_headers: list[tuple[str, str]],
+    duration_attrs: dict[str, str | None] | None = None,
+    sem_conv_opt_in_mode: _StabilityMode = _StabilityMode.DEFAULT,
 ):  # pylint: disable=unused-argument
     """Adds HTTP response attributes to span using the arguments
     passed to a PEP3333-conforming start_response callable.
@@ -501,7 +517,7 @@ def add_response_attributes(
     )
 
 
-def get_default_span_name(environ):
+def get_default_span_name(environ: WSGIEnvironment) -> str:
     """
     Default span name is the HTTP method and URL path, or just the method.
     https://github.com/open-telemetry/opentelemetry-specification/pull/3165
@@ -512,10 +528,12 @@ def get_default_span_name(environ):
     Returns:
         The span name.
     """
-    method = sanitize_method(environ.get("REQUEST_METHOD", "").strip())
+    method = sanitize_method(
+        cast(str, environ.get("REQUEST_METHOD", "")).strip()
+    )
     if method == "_OTHER":
         return "HTTP"
-    path = environ.get("PATH_INFO", "").strip()
+    path = cast(str, environ.get("PATH_INFO", "")).strip()
     if method and path:
         return f"{method} {path}"
     return method
@@ -542,11 +560,11 @@ class OpenTelemetryMiddleware:
 
     def __init__(
         self,
-        wsgi,
-        request_hook=None,
-        response_hook=None,
-        tracer_provider=None,
-        meter_provider=None,
+        wsgi: WSGIApplication,
+        request_hook: RequestHook | None = None,
+        response_hook: ResponseHook | None = None,
+        tracer_provider: TracerProvider | None = None,
+        meter_provider: MeterProvider | None = None,
     ):
         # initialize semantic conventions opt-in if needed
         _OpenTelemetrySemanticConventionStability._initialize()
@@ -593,14 +611,19 @@ def __init__(
 
     @staticmethod
     def _create_start_response(
-        span,
-        start_response,
-        response_hook,
-        duration_attrs,
-        sem_conv_opt_in_mode,
+        span: trace.Span,
+        start_response: StartResponse,
+        response_hook: Callable[[str, list[tuple[str, str]]], None] | None,
+        duration_attrs: dict[str, str | None],
+        sem_conv_opt_in_mode: _StabilityMode,
     ):
         @functools.wraps(start_response)
-        def _start_response(status, response_headers, *args, **kwargs):
+        def _start_response(
+            status: str,
+            response_headers: list[tuple[str, str]],
+            *args: Any,
+            **kwargs: Any,
+        ):
             add_response_attributes(
                 span,
                 status,
@@ -621,7 +644,9 @@ def _start_response(status, response_headers, *args, **kwargs):
         return _start_response
 
     # pylint: disable=too-many-branches
-    def __call__(self, environ, start_response):
+    def __call__(
+        self, environ: WSGIEnvironment, start_response: StartResponse
+    ):
         """The WSGI application
 
         Args:
@@ -703,7 +728,9 @@ def __call__(self, environ, start_response):
 # Put this in a subfunction to not delay the call to the wrapped
 # WSGI application (instrumentation should change the application
 # behavior as little as possible).
-def _end_span_after_iterating(iterable, span, token):
+def _end_span_after_iterating(
+    iterable: Iterable[T], span: trace.Span, token: object
+) -> Iterable[T]:
     try:
         with trace.use_span(span):
             yield from iterable
@@ -717,10 +744,8 @@ def _end_span_after_iterating(iterable, span, token):
 
 
 # TODO: inherit from opentelemetry.instrumentation.propagators.Setter
-
-
 class ResponsePropagationSetter:
-    def set(self, carrier, key, value):  # pylint: disable=no-self-use
+    def set(self, carrier: list[tuple[str, T]], key: str, value: T):  # pylint: disable=no-self-use
         carrier.append((key, value))
 
 

instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/package.py

@@ -12,8 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
 
-_instruments = tuple()
+_instruments: tuple[str, ...] = tuple()
 
 _supports_metrics = True
 

util/opentelemetry-util-http/src/opentelemetry/util/http/__init__.py

@@ -19,7 +19,7 @@
 from re import IGNORECASE as RE_IGNORECASE
 from re import compile as re_compile
 from re import search
-from typing import Callable, Iterable, Optional
+from typing import Callable, Iterable, Optional, overload
 from urllib.parse import urlparse, urlunparse
 
 from opentelemetry.semconv.trace import SpanAttributes
@@ -193,6 +193,14 @@ def normalise_response_header_name(header: str) -> str:
     return f"http.response.header.{key}"
 
 
+@overload
+def sanitize_method(method: str) -> str: ...
+
+
+@overload
+def sanitize_method(method: None) -> None: ...
+
+
 def sanitize_method(method: Optional[str]) -> Optional[str]:
     if method is None:
         return None