From 9b6a71898e2df828e4707d9f1c6d086040b70d72 Mon Sep 17 00:00:00 2001
From: Daniel Szoke <daniel.szoke@sentry.io>
Date: Wed, 3 Jul 2024 17:18:23 +0200
Subject: [PATCH] ref(transport): Improve event data category typing

Done to more clearly define event data categories, in preparation for https://github.com/getsentry/sentry-python/issues/3229.
---
 sentry_sdk/transport.py | 14 ++++++--------
 sentry_sdk/types.py     |  5 +++--
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/sentry_sdk/transport.py b/sentry_sdk/transport.py
index 2cbba041a6..293dfc0e97 100644
--- a/sentry_sdk/transport.py
+++ b/sentry_sdk/transport.py
@@ -34,9 +34,7 @@
     from urllib3.poolmanager import PoolManager
     from urllib3.poolmanager import ProxyManager
 
-    from sentry_sdk._types import Event
-
-    DataCategory = Optional[str]
+    from sentry_sdk._types import Event, EventDataCategory
 
 KEEP_ALIVE_SOCKET_OPTIONS = []
 for option in [
@@ -133,7 +131,7 @@ def kill(self):
     def record_lost_event(
         self,
         reason,  # type: str
-        data_category=None,  # type: Optional[str]
+        data_category=None,  # type: Optional[EventDataCategory]
         item=None,  # type: Optional[Item]
     ):
         # type: (...) -> None
@@ -155,7 +153,7 @@ def __del__(self):
 
 
 def _parse_rate_limits(header, now=None):
-    # type: (Any, Optional[datetime]) -> Iterable[Tuple[DataCategory, datetime]]
+    # type: (Any, Optional[datetime]) -> Iterable[Tuple[Optional[EventDataCategory], datetime]]
     if now is None:
         now = datetime.now(timezone.utc)
 
@@ -195,11 +193,11 @@ def __init__(
         self.options = options  # type: Dict[str, Any]
         self._worker = BackgroundWorker(queue_size=options["transport_queue_size"])
         self._auth = self.parsed_dsn.to_auth("sentry.python/%s" % VERSION)
-        self._disabled_until = {}  # type: Dict[DataCategory, datetime]
+        self._disabled_until = {}  # type: Dict[Optional[EventDataCategory], datetime]
         self._retry = urllib3.util.Retry()
         self._discarded_events = defaultdict(
             int
-        )  # type: DefaultDict[Tuple[str, str], int]
+        )  # type: DefaultDict[Tuple[EventDataCategory, str], int]
         self._last_client_report_sent = time.time()
 
         compresslevel = options.get("_experiments", {}).get(
@@ -224,7 +222,7 @@ def __init__(
     def record_lost_event(
         self,
         reason,  # type: str
-        data_category=None,  # type: Optional[str]
+        data_category=None,  # type: Optional[EventDataCategory]
         item=None,  # type: Optional[Item]
     ):
         # type: (...) -> None
diff --git a/sentry_sdk/types.py b/sentry_sdk/types.py
index 16c57ceea4..a81be8f1c1 100644
--- a/sentry_sdk/types.py
+++ b/sentry_sdk/types.py
@@ -11,13 +11,14 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from sentry_sdk._types import Event, Hint
+    from sentry_sdk._types import Event, EventDataCategory, Hint
 else:
     from typing import Any
 
     # The lines below allow the types to be imported from outside `if TYPE_CHECKING`
     # guards. The types in this module are only intended to be used for type hints.
     Event = Any
+    EventDataCategory = Any
     Hint = Any
 
-__all__ = ("Event", "Hint")
+__all__ = ("Event", "EventDataCategory", "Hint")