Migrate typing for opentelemetry files (#4510)

* remove casts
* generalize `is_valid_sample_rate` to return the validated float
* added a `validate_scopes` helper method to validate the scopes entry
on the context and return the narrowed type
* added a `get_typed_attribute` helper method to narrow the type on
getting attributes from otel
* fix a flaky test on gevent

Author: sl0thentr0py

sentry_sdk/opentelemetry/contextvars_context.py

@@ -1,4 +1,5 @@
-from typing import cast, TYPE_CHECKING
+from __future__ import annotations
+from typing import TYPE_CHECKING
 
 from opentelemetry.trace import get_current_span, set_span_in_context
 from opentelemetry.trace.span import INVALID_SPAN
@@ -13,36 +14,37 @@
     SENTRY_USE_CURRENT_SCOPE_KEY,
     SENTRY_USE_ISOLATION_SCOPE_KEY,
 )
+from sentry_sdk.opentelemetry.scope import PotelScope, validate_scopes
 
 if TYPE_CHECKING:
-    from typing import Optional
     from contextvars import Token
-    import sentry_sdk.opentelemetry.scope as scope
 
 
 class SentryContextVarsRuntimeContext(ContextVarsRuntimeContext):
-    def attach(self, context):
-        # type: (Context) -> Token[Context]
-        scopes = get_value(SENTRY_SCOPES_KEY, context)
+    def attach(self, context: Context) -> Token[Context]:
+        scopes = validate_scopes(get_value(SENTRY_SCOPES_KEY, context))
 
-        should_fork_isolation_scope = context.pop(
-            SENTRY_FORK_ISOLATION_SCOPE_KEY, False
+        should_fork_isolation_scope = bool(
+            context.pop(SENTRY_FORK_ISOLATION_SCOPE_KEY, False)
         )
-        should_fork_isolation_scope = cast("bool", should_fork_isolation_scope)
 
         should_use_isolation_scope = context.pop(SENTRY_USE_ISOLATION_SCOPE_KEY, None)
-        should_use_isolation_scope = cast(
-            "Optional[scope.PotelScope]", should_use_isolation_scope
+        should_use_isolation_scope = (
+            should_use_isolation_scope
+            if isinstance(should_use_isolation_scope, PotelScope)
+            else None
         )
 
         should_use_current_scope = context.pop(SENTRY_USE_CURRENT_SCOPE_KEY, None)
-        should_use_current_scope = cast(
-            "Optional[scope.PotelScope]", should_use_current_scope
+        should_use_current_scope = (
+            should_use_current_scope
+            if isinstance(should_use_current_scope, PotelScope)
+            else None
         )
 
         if scopes:
-            scopes = cast("tuple[scope.PotelScope, scope.PotelScope]", scopes)
-            (current_scope, isolation_scope) = scopes
+            current_scope = scopes[0]
+            isolation_scope = scopes[1]
         else:
             current_scope = sentry_sdk.get_current_scope()
             isolation_scope = sentry_sdk.get_isolation_scope()

sentry_sdk/opentelemetry/propagator.py

@@ -1,4 +1,4 @@
-from typing import cast
+from __future__ import annotations
 
 from opentelemetry import trace
 from opentelemetry.context import (
@@ -37,21 +37,25 @@
     extract_sentrytrace_data,
     should_propagate_trace,
 )
+from sentry_sdk.opentelemetry.scope import validate_scopes
 
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from typing import Optional, Set
-    import sentry_sdk.opentelemetry.scope as scope
 
 
 class SentryPropagator(TextMapPropagator):
     """
     Propagates tracing headers for Sentry's tracing system in a way OTel understands.
     """
 
-    def extract(self, carrier, context=None, getter=default_getter):
-        # type: (CarrierT, Optional[Context], Getter[CarrierT]) -> Context
+    def extract(
+        self,
+        carrier: CarrierT,
+        context: Optional[Context] = None,
+        getter: Getter[CarrierT] = default_getter,
+    ) -> Context:
         if context is None:
             context = get_current()
 
@@ -93,13 +97,16 @@ def extract(self, carrier, context=None, getter=default_getter):
         modified_context = trace.set_span_in_context(span, context)
         return modified_context
 
-    def inject(self, carrier, context=None, setter=default_setter):
-        # type: (CarrierT, Optional[Context], Setter[CarrierT]) -> None
-        scopes = get_value(SENTRY_SCOPES_KEY, context)
+    def inject(
+        self,
+        carrier: CarrierT,
+        context: Optional[Context] = None,
+        setter: Setter[CarrierT] = default_setter,
+    ) -> None:
+        scopes = validate_scopes(get_value(SENTRY_SCOPES_KEY, context))
         if not scopes:
             return
 
-        scopes = cast("tuple[scope.PotelScope, scope.PotelScope]", scopes)
         (current_scope, _) = scopes
 
         span = current_scope.span
@@ -114,6 +121,5 @@ def inject(self, carrier, context=None, setter=default_setter):
             setter.set(carrier, key, value)
 
     @property
-    def fields(self):
-        # type: () -> Set[str]
+    def fields(self) -> Set[str]:
         return {SENTRY_TRACE_HEADER_NAME, BAGGAGE_HEADER_NAME}

sentry_sdk/opentelemetry/sampler.py

@@ -1,5 +1,5 @@
+from __future__ import annotations
 from decimal import Decimal
-from typing import cast
 
 from opentelemetry import trace
 from opentelemetry.sdk.trace.sampling import Sampler, SamplingResult, Decision
@@ -21,15 +21,16 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any, Optional, Sequence, Union
+    from typing import Any, Optional, Sequence
     from opentelemetry.context import Context
     from opentelemetry.trace import Link, SpanKind
     from opentelemetry.trace.span import SpanContext
     from opentelemetry.util.types import Attributes
 
 
-def get_parent_sampled(parent_context, trace_id):
-    # type: (Optional[SpanContext], int) -> Optional[bool]
+def get_parent_sampled(
+    parent_context: Optional[SpanContext], trace_id: int
+) -> Optional[bool]:
     if parent_context is None:
         return None
 
@@ -54,8 +55,9 @@ def get_parent_sampled(parent_context, trace_id):
     return None
 
 
-def get_parent_sample_rate(parent_context, trace_id):
-    # type: (Optional[SpanContext], int) -> Optional[float]
+def get_parent_sample_rate(
+    parent_context: Optional[SpanContext], trace_id: int
+) -> Optional[float]:
     if parent_context is None:
         return None
 
@@ -74,8 +76,9 @@ def get_parent_sample_rate(parent_context, trace_id):
     return None
 
 
-def get_parent_sample_rand(parent_context, trace_id):
-    # type: (Optional[SpanContext], int) -> Optional[Decimal]
+def get_parent_sample_rand(
+    parent_context: Optional[SpanContext], trace_id: int
+) -> Optional[Decimal]:
     if parent_context is None:
         return None
 
@@ -91,8 +94,12 @@ def get_parent_sample_rand(parent_context, trace_id):
     return None
 
 
-def dropped_result(span_context, attributes, sample_rate=None, sample_rand=None):
-    # type: (SpanContext, Attributes, Optional[float], Optional[Decimal]) -> SamplingResult
+def dropped_result(
+    span_context: SpanContext,
+    attributes: Attributes,
+    sample_rate: Optional[float] = None,
+    sample_rand: Optional[Decimal] = None,
+) -> SamplingResult:
     """
     React to a span getting unsampled and return a DROP SamplingResult.
 
@@ -129,8 +136,12 @@ def dropped_result(span_context, attributes, sample_rate=None, sample_rand=None)
     )
 
 
-def sampled_result(span_context, attributes, sample_rate=None, sample_rand=None):
-    # type: (SpanContext, Attributes, Optional[float], Optional[Decimal]) -> SamplingResult
+def sampled_result(
+    span_context: SpanContext,
+    attributes: Attributes,
+    sample_rate: Optional[float] = None,
+    sample_rand: Optional[Decimal] = None,
+) -> SamplingResult:
     """
     React to a span being sampled and return a sampled SamplingResult.
 
@@ -151,8 +162,12 @@ def sampled_result(span_context, attributes, sample_rate=None, sample_rand=None)
     )
 
 
-def _update_trace_state(span_context, sampled, sample_rate=None, sample_rand=None):
-    # type: (SpanContext, bool, Optional[float], Optional[Decimal]) -> TraceState
+def _update_trace_state(
+    span_context: SpanContext,
+    sampled: bool,
+    sample_rate: Optional[float] = None,
+    sample_rand: Optional[Decimal] = None,
+) -> TraceState:
     trace_state = span_context.trace_state
 
     sampled = "true" if sampled else "false"
@@ -175,15 +190,14 @@ def _update_trace_state(span_context, sampled, sample_rate=None, sample_rand=Non
 class SentrySampler(Sampler):
     def should_sample(
         self,
-        parent_context,  # type: Optional[Context]
-        trace_id,  # type: int
-        name,  # type: str
-        kind=None,  # type: Optional[SpanKind]
-        attributes=None,  # type: Attributes
-        links=None,  # type: Optional[Sequence[Link]]
-        trace_state=None,  # type: Optional[TraceState]
-    ):
-        # type: (...) -> SamplingResult
+        parent_context: Optional[Context],
+        trace_id: int,
+        name: str,
+        kind: Optional[SpanKind] = None,
+        attributes: Attributes = None,
+        links: Optional[Sequence[Link]] = None,
+        trace_state: Optional[TraceState] = None,
+    ) -> SamplingResult:
         client = sentry_sdk.get_client()
 
         parent_span_context = trace.get_current_span(parent_context).get_span_context()
@@ -209,13 +223,12 @@ def should_sample(
             sample_rand = parent_sample_rand
         else:
             # We are the head SDK and we need to generate a new sample_rand
-            sample_rand = cast(Decimal, _generate_sample_rand(str(trace_id), (0, 1)))
+            sample_rand = _generate_sample_rand(str(trace_id), (0, 1))
 
         # Explicit sampled value provided at start_span
-        custom_sampled = cast(
-            "Optional[bool]", attributes.get(SentrySpanAttribute.CUSTOM_SAMPLED)
-        )
-        if custom_sampled is not None:
+        custom_sampled = attributes.get(SentrySpanAttribute.CUSTOM_SAMPLED)
+
+        if custom_sampled is not None and isinstance(custom_sampled, bool):
             if is_root_span:
                 sample_rate = float(custom_sampled)
                 if sample_rate > 0:
@@ -262,7 +275,8 @@ def should_sample(
                 sample_rate_to_propagate = sample_rate
 
         # If the sample rate is invalid, drop the span
-        if not is_valid_sample_rate(sample_rate, source=self.__class__.__name__):
+        sample_rate = is_valid_sample_rate(sample_rate, source=self.__class__.__name__)
+        if sample_rate is None:
             logger.warning(
                 f"[Tracing.Sampler] Discarding {name} because of invalid sample rate."
             )
@@ -275,7 +289,6 @@ def should_sample(
                 sample_rate_to_propagate = sample_rate
 
         # Compare sample_rand to sample_rate to make the final sampling decision
-        sample_rate = float(cast("Union[bool, float, int]", sample_rate))
         sampled = sample_rand < Decimal.from_float(sample_rate)
 
         if sampled:
@@ -307,9 +320,13 @@ def get_description(self) -> str:
         return self.__class__.__name__
 
 
-def create_sampling_context(name, attributes, parent_span_context, trace_id):
-    # type: (str, Attributes, Optional[SpanContext], int) -> dict[str, Any]
-    sampling_context = {
+def create_sampling_context(
+    name: str,
+    attributes: Attributes,
+    parent_span_context: Optional[SpanContext],
+    trace_id: int,
+) -> dict[str, Any]:
+    sampling_context: dict[str, Any] = {
         "transaction_context": {
             "name": name,
             "op": attributes.get(SentrySpanAttribute.OP) if attributes else None,
@@ -318,7 +335,7 @@ def create_sampling_context(name, attributes, parent_span_context, trace_id):
             ),
         },
         "parent_sampled": get_parent_sampled(parent_span_context, trace_id),
-    }  # type: dict[str, Any]
+    }
 
     if attributes is not None:
         sampling_context.update(attributes)