From 4137a8d9db174a2fbd03ce9e44334fbc189d7048 Mon Sep 17 00:00:00 2001
From: Katie Byers <katie.byers@sentry.io>
Date: Wed, 21 Oct 2020 10:19:43 -0700
Subject: [PATCH] feat(tracing): Add types for `traces_sampler` implementation
 (#864)

- Types for the `traces_sampler` itself (the function and its input)
- A new attribute on the `Transaction` class tracking the parent sampling decision separately from the sampling decision of the transaction itself, since part of the `traces_sampler` spec is that there needs to be a difference between an inherited decision and an explicitly set decision.
---
 sentry_sdk/_types.py  | 6 ++++++
 sentry_sdk/consts.py  | 8 +++++++-
 sentry_sdk/tracing.py | 4 +++-
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/sentry_sdk/_types.py b/sentry_sdk/_types.py
index 7b727422a1..95e4ac3ba3 100644
--- a/sentry_sdk/_types.py
+++ b/sentry_sdk/_types.py
@@ -5,6 +5,7 @@
 
 
 if MYPY:
+    from numbers import Real
     from types import TracebackType
     from typing import Any
     from typing import Callable
@@ -12,6 +13,7 @@
     from typing import Optional
     from typing import Tuple
     from typing import Type
+    from typing import Union
     from typing_extensions import Literal
 
     ExcInfo = Tuple[
@@ -24,10 +26,14 @@
     Breadcrumb = Dict[str, Any]
     BreadcrumbHint = Dict[str, Any]
 
+    SamplingContext = Dict[str, Any]
+
     EventProcessor = Callable[[Event, Hint], Optional[Event]]
     ErrorProcessor = Callable[[Event, ExcInfo], Optional[Event]]
     BreadcrumbProcessor = Callable[[Breadcrumb, BreadcrumbHint], Optional[Breadcrumb]]
 
+    TracesSampler = Callable[[SamplingContext], Union[Real, bool]]
+
     # https://github.com/python/mypy/issues/5710
     NotImplementedType = Any
 
diff --git a/sentry_sdk/consts.py b/sentry_sdk/consts.py
index cc200107f6..01cc7568fa 100644
--- a/sentry_sdk/consts.py
+++ b/sentry_sdk/consts.py
@@ -14,7 +14,12 @@
     from sentry_sdk.transport import Transport
     from sentry_sdk.integrations import Integration
 
-    from sentry_sdk._types import Event, EventProcessor, BreadcrumbProcessor
+    from sentry_sdk._types import (
+        BreadcrumbProcessor,
+        Event,
+        EventProcessor,
+        TracesSampler,
+    )
 
     # Experiments are feature flags to enable and disable certain unstable SDK
     # functionality. Changing them from the defaults (`None`) in production
@@ -65,6 +70,7 @@ def __init__(
         ca_certs=None,  # type: Optional[str]
         propagate_traces=True,  # type: bool
         traces_sample_rate=0.0,  # type: float
+        traces_sampler=None,  # type: Optional[TracesSampler]
         auto_enabling_integrations=True,  # type: bool
         _experiments={},  # type: Experiments  # noqa: B006
     ):
diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py
index af256d583e..80b4b377d9 100644
--- a/sentry_sdk/tracing.py
+++ b/sentry_sdk/tracing.py
@@ -449,11 +449,12 @@ def get_trace_context(self):
 
 
 class Transaction(Span):
-    __slots__ = ("name",)
+    __slots__ = ("name", "parent_sampled")
 
     def __init__(
         self,
         name="",  # type: str
+        parent_sampled=None,  # type: Optional[bool]
         **kwargs  # type: Any
     ):
         # type: (...) -> None
@@ -468,6 +469,7 @@ def __init__(
             name = kwargs.pop("transaction")
         Span.__init__(self, **kwargs)
         self.name = name
+        self.parent_sampled = parent_sampled
 
     def __repr__(self):
         # type: () -> str